This is an automated email from the ASF dual-hosted git repository. ethanfeng pushed a commit to branch CELEBORN-347 in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
commit ead2c22cff22d57e7a199a3ba74ef2e71c15d60b Author: Ethan Feng <[email protected]> AuthorDate: Tue Feb 28 16:38:05 2023 +0800 [CELEBORN-346][FLINK] Support read multiple locations. --- .../plugin/flink/network/MessageDecoderExt.java | 4 ++ .../plugin/flink/network/ReadClientHandler.java | 8 ++- .../flink/readclient/FlinkShuffleClientImpl.java | 5 +- .../plugin/flink/readclient/RssBufferStream.java | 68 +++++++++++++++++----- .../plugin/flink/RemoteBufferStreamReader.java | 39 ++++++++----- .../plugin/flink/RemoteShuffleInputGate.java | 12 +--- .../common/network/protocol/EndOfStream.java | 52 +++++++++++++++++ .../celeborn/common/network/protocol/Message.java | 8 ++- .../common/network/server/BufferStreamManager.java | 1 + .../common/network/server/DataPartitionReader.java | 39 +++++++------ 10 files changed, 173 insertions(+), 63 deletions(-) diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/MessageDecoderExt.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/MessageDecoderExt.java index fbee2e4fd..38cecf232 100644 --- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/MessageDecoderExt.java +++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/MessageDecoderExt.java @@ -81,6 +81,10 @@ public class MessageDecoderExt { backlog = in.readInt(); return new BacklogAnnouncement(streamId, backlog); + case END_OF_STREAM: + streamId = in.readLong(); + return new EndOfStream(streamId); + default: throw new IllegalArgumentException("Unexpected message type: " + type); } diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java index 8f768c26b..be82d4574 100644 --- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java +++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java @@ -30,10 +30,10 @@ import org.apache.celeborn.common.network.server.BaseMessageHandler; import org.apache.celeborn.plugin.flink.protocol.ReadData; public class ReadClientHandler extends BaseMessageHandler { - private static Logger logger = LoggerFactory.getLogger(ReadClientHandler.class); - private ConcurrentHashMap<Long, Consumer<RequestMessage>> streamHandlers = + private static final Logger logger = LoggerFactory.getLogger(ReadClientHandler.class); + private final ConcurrentHashMap<Long, Consumer<RequestMessage>> streamHandlers = new ConcurrentHashMap<>(); - private ConcurrentHashMap<Long, TransportClient> streamClients = new ConcurrentHashMap<>(); + private final ConcurrentHashMap<Long, TransportClient> streamClients = new ConcurrentHashMap<>(); public void registerHandler( long streamId, Consumer<RequestMessage> handle, TransportClient client) { @@ -75,6 +75,8 @@ public class ReadClientHandler extends BaseMessageHandler { logger.warn("Unexpected streamId received: {}", streamId); } break; + case END_OF_STREAM: + streamHandlers.get(streamId).accept(msg); case ONE_WAY_MESSAGE: // ignore it. break; diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java index 3de2bb1f7..83736afff 100644 --- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java +++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java @@ -37,8 +37,8 @@ public class FlinkShuffleClientImpl extends ShuffleClientImpl { public static final Logger logger = LoggerFactory.getLogger(FlinkShuffleClientImpl.class); private static volatile FlinkShuffleClientImpl _instance; private static volatile boolean initialized = false; - private FlinkTransportClientFactory flinkTransportClientFactory; - private ReadClientHandler readClientHandler = new ReadClientHandler(); + private final FlinkTransportClientFactory flinkTransportClientFactory; + private final ReadClientHandler readClientHandler = new ReadClientHandler(); public static FlinkShuffleClientImpl get( String driverHost, int port, CelebornConf conf, UserIdentifier userIdentifier) { @@ -98,7 +98,6 @@ public class FlinkShuffleClientImpl extends ShuffleClientImpl { return RssBufferStream.empty(); } else { return RssBufferStream.create( - conf, flinkTransportClientFactory, shuffleKey, fileGroups.partitionGroups.get(partitionId).toArray(new PartitionLocation[0]), diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/RssBufferStream.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/RssBufferStream.java index 16cc03419..b8514a5b8 100644 --- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/RssBufferStream.java +++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/RssBufferStream.java @@ -26,7 +26,6 @@ import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import org.apache.celeborn.common.CelebornConf; import org.apache.celeborn.common.network.client.RpcResponseCallback; import org.apache.celeborn.common.network.client.TransportClient; import org.apache.celeborn.common.network.protocol.*; @@ -35,27 +34,29 @@ import org.apache.celeborn.plugin.flink.network.FlinkTransportClientFactory; public class RssBufferStream { - private static Logger logger = LoggerFactory.getLogger(RssBufferStream.class); - private CelebornConf conf; + private static final Logger logger = LoggerFactory.getLogger(RssBufferStream.class); private FlinkTransportClientFactory clientFactory; private String shuffleKey; private PartitionLocation[] locations; private int subIndexStart; private int subIndexEnd; private TransportClient client; - private int currentLocationIndex = 0; - private long streamId = 0; + private volatile int currentLocationIndex = 0; + private volatile long streamId = 0; + private Supplier<ByteBuf> supplier; + private int initialCredit; + private FlinkShuffleClientImpl mapShuffleClient; + private Consumer<RequestMessage> messageConsumer; + private Consumer<Throwable> failureListener; public RssBufferStream() {} public RssBufferStream( - CelebornConf conf, FlinkTransportClientFactory dataClientFactory, String shuffleKey, PartitionLocation[] locations, int subIndexStart, int subIndexEnd) { - this.conf = conf; this.clientFactory = dataClientFactory; this.shuffleKey = shuffleKey; this.locations = locations; @@ -63,12 +64,20 @@ public class RssBufferStream { this.subIndexEnd = subIndexEnd; } - public void open( + public void initByReader( Supplier<ByteBuf> supplier, int initialCredit, FlinkShuffleClientImpl mapShuffleClient, - Consumer<RequestMessage> messageConsumer) - throws IOException, InterruptedException { + Consumer<RequestMessage> messageConsumer, + Consumer<Throwable> failureListener) { + this.supplier = supplier; + this.initialCredit = initialCredit; + this.mapShuffleClient = mapShuffleClient; + this.messageConsumer = messageConsumer; + this.failureListener = failureListener; + } + + public void open() throws IOException, InterruptedException { this.client = clientFactory.createClient( locations[currentLocationIndex].getHost(), @@ -105,9 +114,8 @@ public class RssBufferStream { .writeAndFlush(addCredit) .addListener( future -> { - if (future.isSuccess()) { + if (!future.isSuccess()) { // Send ReadAddCredit do not expect response. - } else { logger.warn( "Send ReadAddCredit to {} failed, detail {}", this.client.getSocketAddress().toString(), @@ -125,7 +133,6 @@ public class RssBufferStream { } public static RssBufferStream create( - CelebornConf conf, FlinkTransportClientFactory dataClientFactory, String shuffleKey, PartitionLocation[] locations, @@ -135,7 +142,7 @@ public class RssBufferStream { return empty(); } else { return new RssBufferStream( - conf, dataClientFactory, shuffleKey, locations, subIndexStart, subIndexEnd); + dataClientFactory, shuffleKey, locations, subIndexStart, subIndexEnd); } } @@ -146,6 +153,37 @@ public class RssBufferStream { } public void close() { - clientFactory.unregisterSupplier(this.getStreamId()); + freeSupplier(streamId); + } + + private void freeSupplier(long streamId) { + clientFactory.unregisterSupplier(streamId); + } + + public void endStream(long streamId) { + if (this.streamId != streamId) { + logger.error("Unexpected end of stream with {}", streamId); + return; + } + + if (currentLocationIndex >= locations.length) { + return; + } + + currentLocationIndex++; + freeSupplier(streamId); + // release netty thread + Thread thread = + new Thread( + () -> { + try { + this.open(); + } catch (Exception e) { + this.failureListener.accept( + new IOException( + "Read location " + locations[currentLocationIndex] + " failed.", e)); + } + }); + thread.start(); } } diff --git a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java index 12cc0f324..c2e04793f 100644 --- a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java +++ b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java @@ -24,6 +24,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.apache.celeborn.common.network.protocol.BacklogAnnouncement; +import org.apache.celeborn.common.network.protocol.EndOfStream; import org.apache.celeborn.common.network.protocol.ReadAddCredit; import org.apache.celeborn.common.network.protocol.RequestMessage; import org.apache.celeborn.plugin.flink.buffer.CreditListener; @@ -33,20 +34,20 @@ import org.apache.celeborn.plugin.flink.readclient.FlinkShuffleClientImpl; import org.apache.celeborn.plugin.flink.readclient.RssBufferStream; public class RemoteBufferStreamReader extends CreditListener { - private static Logger logger = LoggerFactory.getLogger(RemoteBufferStreamReader.class); - private TransferBufferPool bufferPool; - private FlinkShuffleClientImpl client; - private String applicationId; - private int shuffleId; - private int partitionId; - private int subPartitionIndexStart; - private int subPartitionIndexEnd; + private static final Logger logger = LoggerFactory.getLogger(RemoteBufferStreamReader.class); + private final TransferBufferPool bufferPool; + private final FlinkShuffleClientImpl client; + private final String applicationId; + private final int shuffleId; + private final int partitionId; + private final int subPartitionIndexStart; + private final int subPartitionIndexEnd; private boolean isOpen; - private Consumer<ByteBuf> dataListener; - private Consumer<Throwable> failureListener; + private final Consumer<ByteBuf> dataListener; private RssBufferStream bufferStream; private volatile boolean closed = false; - private Consumer<RequestMessage> messageConsumer; + private final Consumer<RequestMessage> messageConsumer; + private final Consumer<Throwable> failureListener; public RemoteBufferStreamReader( FlinkShuffleClientImpl client, @@ -72,6 +73,8 @@ public class RemoteBufferStreamReader extends CreditListener { dataReceived((ReadData) requestMessage); } else if (requestMessage instanceof BacklogAnnouncement) { backlogReceived(((BacklogAnnouncement) requestMessage).getBacklog()); + } else if (requestMessage instanceof EndOfStream) { + endOfStreamReceived(((EndOfStream) requestMessage)); } }; } @@ -81,8 +84,13 @@ public class RemoteBufferStreamReader extends CreditListener { this.bufferStream = client.readBufferedPartition( applicationId, shuffleId, partitionId, subPartitionIndexStart, subPartitionIndexEnd); - bufferStream.open( - RemoteBufferStreamReader.this::requestBuffer, initialCredit, client, messageConsumer); + bufferStream.initByReader( + RemoteBufferStreamReader.this::requestBuffer, + initialCredit, + client, + messageConsumer, + failureListener); + bufferStream.open(); } catch (InterruptedException e) { throw new RuntimeException("Failed to openStream.", e); } @@ -136,4 +144,9 @@ public class RemoteBufferStreamReader extends CreditListener { } dataListener.accept(readData.getFlinkBuffer()); } + + public void endOfStreamReceived(EndOfStream endOfStream) { + bufferStream.endStream(endOfStream.getStreamId()); + client.getReadClientHandler().removeHandler(endOfStream.getStreamId()); + } } diff --git a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java index 1192125cf..0eecc5a54 100644 --- a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java +++ b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGate.java @@ -124,14 +124,11 @@ public class RemoteShuffleInputGate extends IndexedInputGate { /** Number of pending {@link EndOfData} events to be received. */ private long pendingEndOfDataEvents; /** Max concurrent reader count */ - private int numConcurrentReading = Integer.MAX_VALUE; - /** Keep compatibility with streaming mode. */ - private boolean shouldDrainOnEndOfData = true; - + private final int numConcurrentReading = Integer.MAX_VALUE; /** Data decompressor. */ private final BufferDecompressor bufferDecompressor; - private FlinkShuffleClientImpl shuffleClient; + private final FlinkShuffleClientImpl shuffleClient; public RemoteShuffleInputGate( CelebornConf celebornConf, @@ -351,11 +348,6 @@ public class RemoteShuffleInputGate extends IndexedInputGate { return channelsInfo; } - /** Each one corresponds to a reading channel. */ - public List<RemoteBufferStreamReader> getBufferReaders() { - return bufferReaders; - } - private List<InputChannelInfo> createChannelInfos() { return IntStream.range(0, gateDescriptor.getShuffleDescriptors().length) .mapToObj(i -> new InputChannelInfo(gateIndex, i)) diff --git a/common/src/main/java/org/apache/celeborn/common/network/protocol/EndOfStream.java b/common/src/main/java/org/apache/celeborn/common/network/protocol/EndOfStream.java new file mode 100644 index 000000000..3bebbd3fb --- /dev/null +++ b/common/src/main/java/org/apache/celeborn/common/network/protocol/EndOfStream.java @@ -0,0 +1,52 @@ +/* + * 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.celeborn.common.network.protocol; + +import io.netty.buffer.ByteBuf; + +public class EndOfStream extends RequestMessage { + long streamId; + + public EndOfStream(long streamId) { + this.streamId = streamId; + } + + @Override + public int encodedLength() { + return 8; + } + + @Override + public void encode(ByteBuf buf) { + buf.writeLong(streamId); + } + + @Override + public Type type() { + return Type.END_OF_STREAM; + } + + public long getStreamId() { + return streamId; + } + + public static EndOfStream decode(ByteBuf in) { + long streamId = in.readLong(); + return new EndOfStream(streamId); + } +} diff --git a/common/src/main/java/org/apache/celeborn/common/network/protocol/Message.java b/common/src/main/java/org/apache/celeborn/common/network/protocol/Message.java index d2bbd55a6..f1562044f 100644 --- a/common/src/main/java/org/apache/celeborn/common/network/protocol/Message.java +++ b/common/src/main/java/org/apache/celeborn/common/network/protocol/Message.java @@ -92,7 +92,8 @@ public abstract class Message implements Encodable { READ_ADD_CREDIT(16), READ_DATA(17), OPEN_STREAM_WITH_CREDIT(18), - BACKLOG_ANNOUNCEMENT(19); + BACKLOG_ANNOUNCEMENT(19), + END_OF_STREAM(20); private final byte id; Type(int id) { @@ -158,6 +159,8 @@ public abstract class Message implements Encodable { return OPEN_STREAM_WITH_CREDIT; case 19: return BACKLOG_ANNOUNCEMENT; + case 20: + return END_OF_STREAM; case -1: throw new IllegalArgumentException("User type messages cannot be decoded."); default: @@ -223,6 +226,9 @@ public abstract class Message implements Encodable { case BACKLOG_ANNOUNCEMENT: return BacklogAnnouncement.decode(in); + case END_OF_STREAM: + return EndOfStream.decode(in); + default: throw new IllegalArgumentException("Unexpected message type: " + msgType); } diff --git a/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java b/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java index 9e367776d..f9acfb770 100644 --- a/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java +++ b/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java @@ -384,6 +384,7 @@ public class BufferStreamManager { removeStream(streamId); streams.remove(streamId); servingStreams.remove(streamId); + dataPartitionReader.notifyEndOfStream(streamId); return true; } diff --git a/common/src/main/java/org/apache/celeborn/common/network/server/DataPartitionReader.java b/common/src/main/java/org/apache/celeborn/common/network/server/DataPartitionReader.java index 612712147..06da8657c 100644 --- a/common/src/main/java/org/apache/celeborn/common/network/server/DataPartitionReader.java +++ b/common/src/main/java/org/apache/celeborn/common/network/server/DataPartitionReader.java @@ -34,6 +34,7 @@ import org.slf4j.LoggerFactory; import org.apache.celeborn.common.meta.FileInfo; import org.apache.celeborn.common.network.protocol.BacklogAnnouncement; +import org.apache.celeborn.common.network.protocol.EndOfStream; import org.apache.celeborn.common.network.protocol.ReadData; import org.apache.celeborn.common.network.server.memory.Recycler; import org.apache.celeborn.common.network.server.memory.WrappedDataBuffer; @@ -51,12 +52,12 @@ public class DataPartitionReader implements Comparable<DataPartitionReader> { private int currentDataRegion = -1; private long dataConsumingOffset; private volatile long currentPartitionRemainingBytes; - private FileInfo fileInfo; - private int INDEX_ENTRY_SIZE = 16; - private long streamId; + private final FileInfo fileInfo; + private final int INDEX_ENTRY_SIZE = 16; + private final long streamId; protected final Object lock = new Object(); - private AtomicInteger credits = new AtomicInteger(); + private final AtomicInteger credits = new AtomicInteger(); @GuardedBy("lock") protected final Queue<WrappedDataBuffer> buffersRead = new ArrayDeque<>(); @@ -80,11 +81,11 @@ public class DataPartitionReader implements Comparable<DataPartitionReader> { private FileChannel dataFileChannel; private FileChannel indexFileChannel; - private Channel associatedChannel; + private final Channel associatedChannel; - private Runnable recycleStream; + private final Runnable recycleStream; - private AtomicInteger numInFlightRequests = new AtomicInteger(0); + private final AtomicInteger numInFlightRequests = new AtomicInteger(0); public DataPartitionReader( int startPartitionIndex, @@ -121,11 +122,7 @@ public class DataPartitionReader implements Comparable<DataPartitionReader> { public boolean sendWithCredit(int credit) { int oldCredit = credits.getAndAdd(credit); - if (oldCredit == 0) { - return true; - } - - return false; + return oldCredit == 0; } public synchronized boolean readAndSend(Queue<ByteBuf> bufferQueue, Recycler bufferRecycler) @@ -438,11 +435,17 @@ public class DataPartitionReader implements Comparable<DataPartitionReader> { @Override public String toString() { - final StringBuilder sb = new StringBuilder("DataPartitionReader{"); - sb.append("startPartitionIndex=").append(startPartitionIndex); - sb.append(", endPartitionIndex=").append(endPartitionIndex); - sb.append(", streamId=").append(streamId); - sb.append('}'); - return sb.toString(); + return "DataPartitionReader{" + + "startPartitionIndex=" + + startPartitionIndex + + ", endPartitionIndex=" + + endPartitionIndex + + ", streamId=" + + streamId + + '}'; + } + + public void notifyEndOfStream(long streamId) { + this.associatedChannel.writeAndFlush(new EndOfStream(streamId)); } }
