reswqa commented on code in PR #22316: URL: https://github.com/apache/flink/pull/22316#discussion_r1156840081
########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/api/EndOfSegmentEvent.java: ########## @@ -0,0 +1,76 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.api; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.core.memory.DataInputView; +import org.apache.flink.core.memory.DataOutputView; +import org.apache.flink.runtime.event.RuntimeEvent; + +import java.io.IOException; +import java.util.Objects; + +/** EndOfSegmentEvent is used to notify the downstream switch tiers in Tiered Store shuffle mode. */ +@Internal Review Comment: Why marked this as `Internal`. IMH, all `RuntimeEvent` should be stable and ensure compatibility but `Internal` means it might change across releases. ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/api/EndOfSegmentEvent.java: ########## @@ -0,0 +1,76 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.api; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.core.memory.DataInputView; +import org.apache.flink.core.memory.DataOutputView; +import org.apache.flink.runtime.event.RuntimeEvent; + +import java.io.IOException; +import java.util.Objects; + +/** EndOfSegmentEvent is used to notify the downstream switch tiers in Tiered Store shuffle mode. */ Review Comment: ```suggestion /** {@link EndOfSegmentEvent} is used to notify the downstream switch tiers in Tiered Store shuffle mode. */ ``` ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/api/serialization/EventSerializer.java: ########## @@ -139,6 +142,13 @@ public static ByteBuffer toSerializedEvent(AbstractEvent event) throws IOExcepti buf.putInt(selector.getOutputSubtaskIndex()); buf.flip(); return buf; + } else if (eventClass == EndOfSegmentEvent.class) { + EndOfSegmentEvent endOfSegmentEvent = (EndOfSegmentEvent) event; + ByteBuffer buf = ByteBuffer.allocate(8); Review Comment: ```suggestion ByteBuffer buf = ByteBuffer.allocate(2 * Integer.BYTES); ``` ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/Buffer.java: ########## @@ -289,7 +289,10 @@ enum DataType { * Indicates that this subpartition state is fully recovered (emitted). Further data can be * consumed after unblocking. */ - RECOVERY_COMPLETION(false, true, true, false, false); + RECOVERY_COMPLETION(false, true, true, false, false), + + /** Indicates that this event buffer contains the segment id. */ + SEGMENT_EVENT(false, true, false, false, false); Review Comment: There is absolutely no need for this change to be in a separate commit. ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/SubpartitionReader.java: ########## @@ -0,0 +1,42 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel.BufferAndAvailability; + +import java.util.Optional; + +/** {@link SubpartitionReader} is used to consume data from a subpartition of producer. */ +public interface SubpartitionReader { + + /** Set up the SubpartitionReader. */ + void setup(); + + /** + * Get buffer from a single subpartition according to the input channel. + * + * @param inputChannel indicates the specific subpartition. + * @return the nullable buffer. Review Comment: A description like this is meaningless. We'd better describe when this optional has the value and when it is empty. ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/SubpartitionReader.java: ########## @@ -0,0 +1,42 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel.BufferAndAvailability; + +import java.util.Optional; + +/** {@link SubpartitionReader} is used to consume data from a subpartition of producer. */ +public interface SubpartitionReader { Review Comment: I'd prefer rename this to `TieredStoreSubpartitionReader` as `SubpartitionReader` can easily give people the impression that it is the interface of the `Flink Shuffle Service`. ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/SubpartitionReaderImpl.java: ########## @@ -0,0 +1,107 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.runtime.io.network.api.EndOfSegmentEvent; +import org.apache.flink.runtime.io.network.api.serialization.EventSerializer; +import org.apache.flink.runtime.io.network.buffer.Buffer; +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel.BufferAndAvailability; +import org.apache.flink.util.ExceptionUtils; + +import java.io.IOException; +import java.util.List; +import java.util.Optional; +import java.util.function.Consumer; + +import static org.apache.flink.util.Preconditions.checkNotNull; +import static org.apache.flink.util.Preconditions.checkState; + +/** The implementation of {@link SubpartitionReader} interface. */ +public class SubpartitionReaderImpl implements SubpartitionReader { + + private final TierClientFactory clientFactory; + + private final Consumer<InputChannel> queueChannelReceiver; + + private List<TierClient> clientList; + + private int currentSegmentId = 0; + + public SubpartitionReaderImpl( + TierClientFactory clientFactory, Consumer<InputChannel> queueChannelReceiver) { + this.clientFactory = clientFactory; + this.queueChannelReceiver = queueChannelReceiver; + } + + @Override + public void setup() { + this.clientList = clientFactory.createClientList(); + } + + @Override + public Optional<BufferAndAvailability> getNextBuffer(InputChannel inputChannel) { + Optional<BufferAndAvailability> bufferAndAvailability = Optional.empty(); + for (TierClient client : clientList) { + bufferAndAvailability = client.getNextBuffer(inputChannel, currentSegmentId); + if (bufferAndAvailability.isPresent()) { + break; + } + } + if (!bufferAndAvailability.isPresent()) { + return Optional.empty(); + } + BufferAndAvailability bufferData = bufferAndAvailability.get(); + if (bufferData.buffer().getDataType() == Buffer.DataType.SEGMENT_EVENT) { + checkState( + getSegmentId(bufferData) == (currentSegmentId + 1), Review Comment: ```suggestion getSegmentId(bufferData) == currentSegmentId + 1, ``` Please verify if the brackets can be removed, I can't remember their priority :). ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/TierClient.java: ########## @@ -0,0 +1,32 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; + +import java.util.Optional; + +/** The interface of {@link TierClient} in Tiered Store. */ +public interface TierClient { + + Optional<InputChannel.BufferAndAvailability> getNextBuffer( + InputChannel inputChannel, int segmentId); + + void close(); Review Comment: I see this is a no-op method for `LocalTierClient`, is this used for remote tiered client to release resource? ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/api/serialization/EventSerializer.java: ########## @@ -183,6 +193,8 @@ public static AbstractEvent fromSerializedEvent(ByteBuffer buffer, ClassLoader c return new EventAnnouncement(announcedEvent, sequenceNumber); } else if (type == VIRTUAL_CHANNEL_SELECTOR_EVENT) { return new SubtaskConnectionDescriptor(buffer.getInt(), buffer.getInt()); + } else if (type == END_OF_SEGMENT) { + return deserializeEndOfSegment(buffer); Review Comment: I'd prefer inline this method since this is a very simple logic no need to extract to a standalone method. ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/SubpartitionReaderImpl.java: ########## @@ -0,0 +1,107 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.runtime.io.network.api.EndOfSegmentEvent; +import org.apache.flink.runtime.io.network.api.serialization.EventSerializer; +import org.apache.flink.runtime.io.network.buffer.Buffer; +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel.BufferAndAvailability; +import org.apache.flink.util.ExceptionUtils; + +import java.io.IOException; +import java.util.List; +import java.util.Optional; +import java.util.function.Consumer; + +import static org.apache.flink.util.Preconditions.checkNotNull; +import static org.apache.flink.util.Preconditions.checkState; + +/** The implementation of {@link SubpartitionReader} interface. */ +public class SubpartitionReaderImpl implements SubpartitionReader { + + private final TierClientFactory clientFactory; + + private final Consumer<InputChannel> queueChannelReceiver; + + private List<TierClient> clientList; + + private int currentSegmentId = 0; + + public SubpartitionReaderImpl( + TierClientFactory clientFactory, Consumer<InputChannel> queueChannelReceiver) { + this.clientFactory = clientFactory; + this.queueChannelReceiver = queueChannelReceiver; + } + + @Override + public void setup() { + this.clientList = clientFactory.createClientList(); + } + + @Override + public Optional<BufferAndAvailability> getNextBuffer(InputChannel inputChannel) { + Optional<BufferAndAvailability> bufferAndAvailability = Optional.empty(); + for (TierClient client : clientList) { + bufferAndAvailability = client.getNextBuffer(inputChannel, currentSegmentId); + if (bufferAndAvailability.isPresent()) { + break; + } + } + if (!bufferAndAvailability.isPresent()) { + return Optional.empty(); + } + BufferAndAvailability bufferData = bufferAndAvailability.get(); + if (bufferData.buffer().getDataType() == Buffer.DataType.SEGMENT_EVENT) { + checkState( Review Comment: It seems that the segment id contains in `SEGMENT_EVENT` is only used for this check. I'd prefer do not add this sanity check to avoid extra network overhead in hot-path. This will allows us to make it a singleton. Maybe we could rename `SEGMENT_EVENT` to `ADD_SEGMENT_ID_EVENT` to make it more expressive. ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/TierClient.java: ########## @@ -0,0 +1,32 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; + +import java.util.Optional; + +/** The interface of {@link TierClient} in Tiered Store. */ Review Comment: This java doc is really meaningless, we should explain clearly what role it plays in the network layer. Please keep in mind: Java doc is to help other developers understand our protocol. The clearer the definition of this, the more people will help maintain these codes. ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/LocalTierClient.java: ########## @@ -0,0 +1,60 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.annotation.VisibleForTesting; +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; +import org.apache.flink.util.ExceptionUtils; + +import java.io.IOException; +import java.util.Optional; + +/** + * The {@link LocalTierClient} is used to fetch data from memory tier and disk tier in tiered store. + */ +public class LocalTierClient implements TierClient { + + private int latestSegmentId = 0; + + @Override + public Optional<InputChannel.BufferAndAvailability> getNextBuffer( + InputChannel inputChannel, int segmentId) { + Optional<InputChannel.BufferAndAvailability> buffer = Optional.empty(); + if (segmentId > 0L && (segmentId != latestSegmentId)) { + latestSegmentId = segmentId; + inputChannel.notifyRequiredSegmentId(segmentId); + } + try { + buffer = inputChannel.getNextBuffer(); + } catch (IOException | InterruptedException e) { + ExceptionUtils.rethrow(e, "LocalTierClient failed to get next buffer."); + } + return buffer; Review Comment: ```suggestion if (segmentId > 0 && (segmentId != latestSegmentId)) { latestSegmentId = segmentId; inputChannel.notifyRequiredSegmentId(segmentId); } try { return inputChannel.getNextBuffer(); } catch (IOException | InterruptedException e) { throw new RuntimeException("LocalTierClient failed to get next buffer.", e); } ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/TierClient.java: ########## @@ -0,0 +1,32 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; + +import java.util.Optional; + +/** The interface of {@link TierClient} in Tiered Store. */ +public interface TierClient { + + Optional<InputChannel.BufferAndAvailability> getNextBuffer( Review Comment: All methods in public interface needs java doc. ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/SubpartitionReaderImpl.java: ########## @@ -0,0 +1,107 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.runtime.io.network.api.EndOfSegmentEvent; +import org.apache.flink.runtime.io.network.api.serialization.EventSerializer; +import org.apache.flink.runtime.io.network.buffer.Buffer; +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel.BufferAndAvailability; +import org.apache.flink.util.ExceptionUtils; + +import java.io.IOException; +import java.util.List; +import java.util.Optional; +import java.util.function.Consumer; + +import static org.apache.flink.util.Preconditions.checkNotNull; +import static org.apache.flink.util.Preconditions.checkState; + +/** The implementation of {@link SubpartitionReader} interface. */ +public class SubpartitionReaderImpl implements SubpartitionReader { + + private final TierClientFactory clientFactory; + + private final Consumer<InputChannel> queueChannelReceiver; + + private List<TierClient> clientList; + + private int currentSegmentId = 0; + + public SubpartitionReaderImpl( + TierClientFactory clientFactory, Consumer<InputChannel> queueChannelReceiver) { + this.clientFactory = clientFactory; + this.queueChannelReceiver = queueChannelReceiver; + } + + @Override + public void setup() { + this.clientList = clientFactory.createClientList(); + } + + @Override + public Optional<BufferAndAvailability> getNextBuffer(InputChannel inputChannel) { + Optional<BufferAndAvailability> bufferAndAvailability = Optional.empty(); + for (TierClient client : clientList) { + bufferAndAvailability = client.getNextBuffer(inputChannel, currentSegmentId); + if (bufferAndAvailability.isPresent()) { + break; + } + } + if (!bufferAndAvailability.isPresent()) { + return Optional.empty(); + } + BufferAndAvailability bufferData = bufferAndAvailability.get(); + if (bufferData.buffer().getDataType() == Buffer.DataType.SEGMENT_EVENT) { + checkState( + getSegmentId(bufferData) == (currentSegmentId + 1), + "Received illegal segmentId."); + currentSegmentId++; + bufferData.buffer().recycleBuffer(); + queueChannelReceiver.accept(inputChannel); + return getNextBuffer(inputChannel); + } + return Optional.of(bufferData); + } + + @Override + public void close() { + for (TierClient client : clientList) { Review Comment: `close` It is not guaranteed to be called only by the main thread. `clientList` may not thread safe which can cause `NPE`. ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/TieredStoreReader.java: ########## @@ -0,0 +1,48 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel.BufferAndAvailability; + +import java.util.Optional; +import java.util.function.Consumer; + +/** The interface of {@link TieredStoreReader} in Tiered Store. */ +public interface TieredStoreReader { + + /** + * Set up the TieredStoreReader. + * + * @param inputChannels indicate the all producers it consumes. + * @param queueChannelReceiver receives a channel and enqueue it. + */ + void setup(InputChannel[] inputChannels, Consumer<InputChannel> queueChannelReceiver); Review Comment: UUIC, This `Consumer<InputChannel>` can be passing by constructor instead of lazy `setup`. ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/SubpartitionReader.java: ########## @@ -0,0 +1,42 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel.BufferAndAvailability; + +import java.util.Optional; + +/** {@link SubpartitionReader} is used to consume data from a subpartition of producer. */ +public interface SubpartitionReader { + + /** Set up the SubpartitionReader. */ + void setup(); + + /** + * Get buffer from a single subpartition according to the input channel. + * + * @param inputChannel indicates the specific subpartition. + * @return the nullable buffer. + */ + Optional<BufferAndAvailability> getNextBuffer(InputChannel inputChannel); + + /** Close the subpartition reader. */ + void close(); Review Comment: We'd better implements `AutoClosable` here. ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/netty/NettyMessage.java: ########## @@ -895,6 +898,52 @@ public String toString() { } } + /** Message to notify producer about the required segment id. */ + static class SegmentId extends NettyMessage { + + private static final byte ID = 11; + + final int segmentId; + + final InputChannelID receiverId; + + SegmentId(int segmentId, InputChannelID receiverId) { + checkArgument(segmentId > 0L, "The segmentId should be greater than 0"); Review Comment: ```suggestion checkArgument(segmentId > 0, "Non-positive segmentId is illegal."); ``` ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannel.java: ########## @@ -114,7 +114,7 @@ public void checkpointStopped(long checkpointId) { } @Override - protected void requestSubpartition() throws IOException { + public void requestSubpartition() throws IOException { Review Comment: Why change this modifier? ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/SingleInputGate.java: ########## @@ -774,11 +785,13 @@ private Optional<InputWithData<InputChannel, BufferAndAvailability>> waitAndGetN if (!inputChannelOpt.isPresent()) { return Optional.empty(); } - final InputChannel inputChannel = inputChannelOpt.get(); - Optional<BufferAndAvailability> bufferAndAvailabilityOpt = - inputChannel.getNextBuffer(); - + Optional<BufferAndAvailability> bufferAndAvailabilityOpt; + if (tieredStoreReader != null) { + bufferAndAvailabilityOpt = tieredStoreReader.getNextBuffer(inputChannel); + } else { + bufferAndAvailabilityOpt = inputChannel.getNextBuffer(); Review Comment: Here it should be possible to unify ordinary read and tired store read into one interface or encapsulation. ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/LocalTierClient.java: ########## @@ -0,0 +1,60 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.annotation.VisibleForTesting; +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; +import org.apache.flink.util.ExceptionUtils; + +import java.io.IOException; +import java.util.Optional; + +/** + * The {@link LocalTierClient} is used to fetch data from memory tier and disk tier in tiered store. + */ +public class LocalTierClient implements TierClient { + + private int latestSegmentId = 0; + + @Override + public Optional<InputChannel.BufferAndAvailability> getNextBuffer( + InputChannel inputChannel, int segmentId) { + Optional<InputChannel.BufferAndAvailability> buffer = Optional.empty(); + if (segmentId > 0L && (segmentId != latestSegmentId)) { + latestSegmentId = segmentId; + inputChannel.notifyRequiredSegmentId(segmentId); + } + try { + buffer = inputChannel.getNextBuffer(); + } catch (IOException | InterruptedException e) { + ExceptionUtils.rethrow(e, "LocalTierClient failed to get next buffer."); + } + return buffer; + } + + @Override + public void close() { + // nothing to do. + } + + @VisibleForTesting + public int getLatestSegmentId() { Review Comment: I'd prefer do not introduce this `VisibleForTesting` method. We can verify this by intercepting `inputChannel.notifyRequiredSegmentId(segmentId)`. ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/LocalTierClient.java: ########## @@ -0,0 +1,60 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.annotation.VisibleForTesting; +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; +import org.apache.flink.util.ExceptionUtils; + +import java.io.IOException; +import java.util.Optional; + +/** + * The {@link LocalTierClient} is used to fetch data from memory tier and disk tier in tiered store. + */ +public class LocalTierClient implements TierClient { + + private int latestSegmentId = 0; + + @Override + public Optional<InputChannel.BufferAndAvailability> getNextBuffer( + InputChannel inputChannel, int segmentId) { + Optional<InputChannel.BufferAndAvailability> buffer = Optional.empty(); + if (segmentId > 0L && (segmentId != latestSegmentId)) { Review Comment: ```suggestion if (segmentId > 0 && (segmentId != latestSegmentId)) { ``` ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/TieredStoreReaderImpl.java: ########## @@ -0,0 +1,59 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; + +import java.util.Optional; +import java.util.function.Consumer; + +/** The implementation of {@link TieredStoreReader} interface. */ +public class TieredStoreReaderImpl implements TieredStoreReader { + + private final SubpartitionReader[] subpartitionReaders; + + private final int numInputChannels; + + public TieredStoreReaderImpl(int numInputChannels) { + this.numInputChannels = numInputChannels; + this.subpartitionReaders = new SubpartitionReader[numInputChannels]; + } + + @Override + public void setup(InputChannel[] channels, Consumer<InputChannel> queueChannelReceiver) { Review Comment: Why `channels` is needed here? ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/Buffer.java: ########## @@ -289,7 +289,10 @@ enum DataType { * Indicates that this subpartition state is fully recovered (emitted). Further data can be * consumed after unblocking. */ - RECOVERY_COMPLETION(false, true, true, false, false); + RECOVERY_COMPLETION(false, true, true, false, false), + + /** Indicates that this event buffer contains the segment id. */ Review Comment: This java doc does not explain the purpose of introducing this `DataType` at all, please refer to others. ########## flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/tieredstore/downstream/TieredStoreReaderImpl.java: ########## @@ -0,0 +1,59 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.runtime.io.network.partition.tieredstore.downstream; + +import org.apache.flink.runtime.io.network.partition.consumer.InputChannel; + +import java.util.Optional; +import java.util.function.Consumer; + +/** The implementation of {@link TieredStoreReader} interface. */ +public class TieredStoreReaderImpl implements TieredStoreReader { + + private final SubpartitionReader[] subpartitionReaders; + + private final int numInputChannels; + + public TieredStoreReaderImpl(int numInputChannels) { + this.numInputChannels = numInputChannels; + this.subpartitionReaders = new SubpartitionReader[numInputChannels]; + } + + @Override + public void setup(InputChannel[] channels, Consumer<InputChannel> queueChannelReceiver) { + TierClientFactory clientFactory = new TierClientFactory(); Review Comment: I feel a little strange: why do we introduce a `Factory` here instead of directly creating the corresponding object or passing it in from the outside? -- 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. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
