This is an automated email from the ASF dual-hosted git repository. zhangshenghang pushed a commit to branch fix-socket-source-reader-lifecycle in repository https://gitbox.apache.org/repos/asf/seatunnel.git
commit 8c6f33c72351aca8bc34f9287c03f2b66e9b628a Author: zhangshenghang <[email protected]> AuthorDate: Fri Jun 19 22:49:33 2026 +0800 [Fix][Connector-V2] Keep socket reader stream open --- .../socket/source/SocketSourceReader.java | 43 ++++----- .../socket/source/SocketSourceReaderTest.java | 104 +++++++++++++++++++++ 2 files changed, 126 insertions(+), 21 deletions(-) diff --git a/seatunnel-connectors-v2/connector-socket/src/main/java/org/apache/seatunnel/connectors/seatunnel/socket/source/SocketSourceReader.java b/seatunnel-connectors-v2/connector-socket/src/main/java/org/apache/seatunnel/connectors/seatunnel/socket/source/SocketSourceReader.java index 70b4326f1c..4ecaaf4baa 100644 --- a/seatunnel-connectors-v2/connector-socket/src/main/java/org/apache/seatunnel/connectors/seatunnel/socket/source/SocketSourceReader.java +++ b/seatunnel-connectors-v2/connector-socket/src/main/java/org/apache/seatunnel/connectors/seatunnel/socket/source/SocketSourceReader.java @@ -38,6 +38,7 @@ public class SocketSourceReader extends AbstractSingleSplitReader<SeaTunnelRow> private final SocketConfig parameter; private final SingleSplitReaderContext context; private Socket socket; + private BufferedReader reader; private final String delimiter = "\n"; SocketSourceReader(SocketConfig parameter, SingleSplitReaderContext context) { @@ -54,11 +55,14 @@ public class SocketSourceReader extends AbstractSingleSplitReader<SeaTunnelRow> this.parameter.getPort()); socket.connect( new InetSocketAddress(this.parameter.getHost(), this.parameter.getPort()), 0); + reader = new BufferedReader(new InputStreamReader(socket.getInputStream())); } @Override public void close() throws IOException { - if (socket != null) { + if (reader != null) { + reader.close(); + } else if (socket != null) { socket.close(); } } @@ -66,28 +70,25 @@ public class SocketSourceReader extends AbstractSingleSplitReader<SeaTunnelRow> @Override public void pollNext(Collector<SeaTunnelRow> output) throws Exception { StringBuilder buffer = new StringBuilder(); - try (BufferedReader reader = - new BufferedReader(new InputStreamReader(socket.getInputStream()))) { - char[] buf = new char[CHAR_BUFFER_SIZE]; - int bytesRead; - while ((bytesRead = reader.read(buf)) != -1) { - buffer.append(buf, 0, bytesRead); + char[] buf = new char[CHAR_BUFFER_SIZE]; + int bytesRead; + while ((bytesRead = reader.read(buf)) != -1) { + buffer.append(buf, 0, bytesRead); - int delimPos; - while (buffer.length() >= this.delimiter.length() - && (delimPos = buffer.indexOf(this.delimiter)) != -1) { - String record = buffer.substring(0, delimPos); - if (record.endsWith("\r")) { - record = record.substring(0, record.length() - 1); - } - output.collect(new SeaTunnelRow(new Object[] {record})); - buffer.delete(0, delimPos + this.delimiter.length()); - } - if (Boundedness.BOUNDED.equals(context.getBoundedness())) { - // signal to the source that we have reached the end of the data. - context.signalNoMoreElement(); - break; + int delimPos; + while (buffer.length() >= this.delimiter.length() + && (delimPos = buffer.indexOf(this.delimiter)) != -1) { + String record = buffer.substring(0, delimPos); + if (record.endsWith("\r")) { + record = record.substring(0, record.length() - 1); } + output.collect(new SeaTunnelRow(new Object[] {record})); + buffer.delete(0, delimPos + this.delimiter.length()); + } + if (Boundedness.BOUNDED.equals(context.getBoundedness())) { + // signal to the source that we have reached the end of the data. + context.signalNoMoreElement(); + break; } } if (buffer.length() > 0) { diff --git a/seatunnel-connectors-v2/connector-socket/src/test/java/org/apache/seatunnel/connectors/seatunnel/socket/source/SocketSourceReaderTest.java b/seatunnel-connectors-v2/connector-socket/src/test/java/org/apache/seatunnel/connectors/seatunnel/socket/source/SocketSourceReaderTest.java new file mode 100644 index 0000000000..257dfc0b33 --- /dev/null +++ b/seatunnel-connectors-v2/connector-socket/src/test/java/org/apache/seatunnel/connectors/seatunnel/socket/source/SocketSourceReaderTest.java @@ -0,0 +1,104 @@ +/* + * 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.seatunnel.connectors.seatunnel.socket.source; + +import org.apache.seatunnel.api.configuration.ReadonlyConfig; +import org.apache.seatunnel.api.source.Boundedness; +import org.apache.seatunnel.api.source.Collector; +import org.apache.seatunnel.api.source.SourceReader; +import org.apache.seatunnel.api.table.type.SeaTunnelRow; +import org.apache.seatunnel.connectors.seatunnel.common.source.SingleSplitReaderContext; +import org.apache.seatunnel.connectors.seatunnel.socket.config.SocketCommonOptions; +import org.apache.seatunnel.connectors.seatunnel.socket.config.SocketConfig; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.io.BufferedWriter; +import java.io.IOException; +import java.io.OutputStreamWriter; +import java.net.ServerSocket; +import java.net.Socket; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +class SocketSourceReaderTest { + + @Test + void shouldKeepSocketOpenAcrossPollNextCalls() throws Exception { + try (ServerSocket serverSocket = new ServerSocket(0)) { + SocketSourceReader reader = createReader(serverSocket.getLocalPort()); + reader.open(); + try (Socket socket = serverSocket.accept(); + BufferedWriter writer = + new BufferedWriter( + new OutputStreamWriter( + socket.getOutputStream(), StandardCharsets.UTF_8))) { + TestCollector collector = new TestCollector(); + + writeLine(writer, "first"); + reader.pollNext(collector); + writeLine(writer, "second"); + reader.pollNext(collector); + + Assertions.assertEquals("first", collector.rows.get(0).getField(0)); + Assertions.assertEquals("second", collector.rows.get(1).getField(0)); + } finally { + reader.close(); + } + } + } + + private static SocketSourceReader createReader(int port) { + Map<String, Object> options = new HashMap<>(); + options.put(SocketCommonOptions.HOST.key(), "127.0.0.1"); + options.put(SocketCommonOptions.PORT.key(), port); + SourceReader.Context context = mock(SourceReader.Context.class); + when(context.getBoundedness()).thenReturn(Boundedness.BOUNDED); + return new SocketSourceReader( + new SocketConfig(ReadonlyConfig.fromMap(options)), + new SingleSplitReaderContext(context)); + } + + private static void writeLine(BufferedWriter writer, String value) throws IOException { + writer.write(value); + writer.newLine(); + writer.flush(); + } + + private static class TestCollector implements Collector<SeaTunnelRow> { + private final Object lock = new Object(); + private final List<SeaTunnelRow> rows = new ArrayList<>(); + + @Override + public void collect(SeaTunnelRow row) { + rows.add(row); + } + + @Override + public Object getCheckpointLock() { + return lock; + } + } +}
