Skip to content

Commit c327b5f

Browse files
nongliwesm
nongli
authored and
wesm
committed
ARROW-506: Java: Implement echo server for integration testing.
While implementing this, it became clear it made sense for the stream writer to have an API to indicate EOS without closing the stream. The current message the reader is expecting is a 4 byte size for the next batch. This patch proposes we allow 0 as the size to indicate EOS. Author: Nong Li <nongli@gmail.com> Closes apache#295 from nongli/echo_server and squashes the following commits: c115b02 [Nong Li] Add license header. a3a50ca [Nong Li] ARROW-506: Java: Implement echo server for integration testing.
1 parent 69cdbd8 commit c327b5f

7 files changed

Lines changed: 278 additions & 12 deletions

File tree

‎java/tools/src/main/java/org/apache/arrow/tools/EchoServer.java‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
/**
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.arrow.tools;
19+
20+
import java.io.IOException;
21+
import java.io.InputStream;
22+
import java.io.OutputStream;
23+
import java.net.ServerSocket;
24+
import java.net.Socket;
25+
import java.util.ArrayList;
26+
import java.util.List;
27+
28+
import org.apache.arrow.memory.BufferAllocator;
29+
import org.apache.arrow.memory.RootAllocator;
30+
import org.apache.arrow.vector.schema.ArrowRecordBatch;
31+
import org.apache.arrow.vector.stream.ArrowStreamReader;
32+
import org.apache.arrow.vector.stream.ArrowStreamWriter;
33+
import org.slf4j.Logger;
34+
import org.slf4j.LoggerFactory;
35+
36+
import com.google.common.base.Preconditions;
37+
38+
public class EchoServer {
39+
private static final Logger LOGGER = LoggerFactory.getLogger(EchoServer.class);
40+
41+
private boolean closed = false;
42+
private final ServerSocket serverSocket;
43+
44+
public EchoServer(int port) throws IOException {
45+
LOGGER.info("Starting echo server.");
46+
serverSocket = new ServerSocket(port);
47+
LOGGER.info("Running echo server on port: " + port());
48+
}
49+
50+
public int port() { return serverSocket.getLocalPort(); }
51+
52+
public static class ClientConnection implements AutoCloseable {
53+
public final Socket socket;
54+
public ClientConnection(Socket socket) {
55+
this.socket = socket;
56+
}
57+
58+
public void run() throws IOException {
59+
BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
60+
List<ArrowRecordBatch> batches = new ArrayList<ArrowRecordBatch>();
61+
try (
62+
InputStream in = socket.getInputStream();
63+
OutputStream out = socket.getOutputStream();
64+
ArrowStreamReader reader = new ArrowStreamReader(in, allocator);
65+
) {
66+
// Read the entire input stream.
67+
reader.init();
68+
while (true) {
69+
ArrowRecordBatch batch = reader.nextRecordBatch();
70+
if (batch == null) break;
71+
batches.add(batch);
72+
}
73+
LOGGER.info(String.format("Received %d batches", batches.size()));
74+
75+
// Write it back
76+
try (ArrowStreamWriter writer = new ArrowStreamWriter(out, reader.getSchema())) {
77+
for (ArrowRecordBatch batch: batches) {
78+
writer.writeRecordBatch(batch);
79+
}
80+
writer.end();
81+
Preconditions.checkState(reader.bytesRead() == writer.bytesWritten());
82+
}
83+
LOGGER.info("Done writing stream back.");
84+
}
85+
}
86+
87+
@Override
88+
public void close() throws IOException {
89+
socket.close();
90+
}
91+
}
92+
93+
public void run() throws IOException {
94+
try {
95+
while (!closed) {
96+
LOGGER.info("Waiting to accept new client connection.");
97+
Socket clientSocket = serverSocket.accept();
98+
LOGGER.info("Accepted new client connection.");
99+
try (ClientConnection client = new ClientConnection(clientSocket)) {
100+
try {
101+
client.run();
102+
} catch (IOException e) {
103+
LOGGER.warn("Error handling client connection.", e);
104+
}
105+
}
106+
LOGGER.info("Closed connection with client");
107+
}
108+
} catch (java.net.SocketException ex) {
109+
if (!closed) throw ex;
110+
} finally {
111+
serverSocket.close();
112+
LOGGER.info("Server closed.");
113+
}
114+
}
115+
116+
public void close() throws IOException {
117+
closed = true;
118+
serverSocket.close();
119+
}
120+
121+
public static void main(String[] args) throws Exception {
122+
int port;
123+
if (args.length > 0) {
124+
port = Integer.parseInt(args[0]);
125+
} else {
126+
port = 8080;
127+
}
128+
new EchoServer(port).run();
129+
}
130+
}

‎java/tools/src/test/java/org/apache/arrow/tools/EchoServerTest.java‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
/**
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.arrow.tools;
19+
20+
import static java.util.Arrays.asList;
21+
import static org.junit.Assert.assertEquals;
22+
import static org.junit.Assert.assertTrue;
23+
24+
import java.io.IOException;
25+
import java.net.Socket;
26+
import java.net.UnknownHostException;
27+
import java.util.ArrayList;
28+
import java.util.Collections;
29+
import java.util.List;
30+
31+
import org.apache.arrow.memory.BufferAllocator;
32+
import org.apache.arrow.memory.RootAllocator;
33+
import org.apache.arrow.vector.schema.ArrowFieldNode;
34+
import org.apache.arrow.vector.schema.ArrowRecordBatch;
35+
import org.apache.arrow.vector.stream.ArrowStreamReader;
36+
import org.apache.arrow.vector.stream.ArrowStreamWriter;
37+
import org.apache.arrow.vector.types.pojo.ArrowType;
38+
import org.apache.arrow.vector.types.pojo.Field;
39+
import org.apache.arrow.vector.types.pojo.Schema;
40+
import org.junit.Test;
41+
42+
import io.netty.buffer.ArrowBuf;
43+
44+
public class EchoServerTest {
45+
public static ArrowBuf buf(BufferAllocator alloc, byte[] bytes) {
46+
ArrowBuf buffer = alloc.buffer(bytes.length);
47+
buffer.writeBytes(bytes);
48+
return buffer;
49+
}
50+
51+
public static byte[] array(ArrowBuf buf) {
52+
byte[] bytes = new byte[buf.readableBytes()];
53+
buf.readBytes(bytes);
54+
return bytes;
55+
}
56+
57+
private void testEchoServer(int serverPort, Schema schema, List<ArrowRecordBatch> batches)
58+
throws UnknownHostException, IOException {
59+
BufferAllocator alloc = new RootAllocator(Long.MAX_VALUE);
60+
try (Socket socket = new Socket("localhost", serverPort);
61+
ArrowStreamWriter writer = new ArrowStreamWriter(socket.getOutputStream(), schema);
62+
ArrowStreamReader reader = new ArrowStreamReader(socket.getInputStream(), alloc)) {
63+
for (ArrowRecordBatch batch: batches) {
64+
writer.writeRecordBatch(batch);
65+
}
66+
writer.end();
67+
68+
reader.init();
69+
assertEquals(schema, reader.getSchema());
70+
for (int i = 0; i < batches.size(); i++) {
71+
ArrowRecordBatch result = reader.nextRecordBatch();
72+
ArrowRecordBatch expected = batches.get(i);
73+
assertTrue(result != null);
74+
assertEquals(expected.getBuffers().size(), result.getBuffers().size());
75+
for (int j = 0; j < expected.getBuffers().size(); j++) {
76+
assertTrue(expected.getBuffers().get(j).compareTo(result.getBuffers().get(j)) == 0);
77+
}
78+
}
79+
ArrowRecordBatch result = reader.nextRecordBatch();
80+
assertTrue(result == null);
81+
assertEquals(reader.bytesRead(), writer.bytesWritten());
82+
}
83+
}
84+
85+
@Test
86+
public void basicTest() throws InterruptedException, IOException {
87+
final EchoServer server = new EchoServer(0);
88+
int serverPort = server.port();
89+
Thread serverThread = new Thread() {
90+
@Override
91+
public void run() {
92+
try {
93+
server.run();
94+
} catch (IOException e) {
95+
e.printStackTrace();
96+
}
97+
}
98+
};
99+
serverThread.start();
100+
101+
BufferAllocator alloc = new RootAllocator(Long.MAX_VALUE);
102+
byte[] validity = new byte[] { (byte)255, 0};
103+
byte[] values = new byte[] { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16};
104+
ArrowBuf validityb = buf(alloc, validity);
105+
ArrowBuf valuesb = buf(alloc, values);
106+
ArrowRecordBatch batch = new ArrowRecordBatch(
107+
16, asList(new ArrowFieldNode(16, 8)), asList(validityb, valuesb));
108+
109+
Schema schema = new Schema(asList(new Field(
110+
"testField", true, new ArrowType.Int(8, true), Collections.<Field>emptyList())));
111+
112+
// Try an empty stream, just the header.
113+
testEchoServer(serverPort, schema, new ArrayList<ArrowRecordBatch>());
114+
115+
// Try with one batch.
116+
List<ArrowRecordBatch> batches = new ArrayList<>();
117+
batches.add(batch);
118+
testEchoServer(serverPort, schema, batches);
119+
120+
// Try with a few
121+
for (int i = 0; i < 10; i++) {
122+
batches.add(batch);
123+
}
124+
testEchoServer(serverPort, schema, batches);
125+
126+
server.close();
127+
serverThread.join();
128+
}
129+
}

‎java/vector/src/main/java/org/apache/arrow/vector/stream/ArrowStreamWriter.java‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,14 +35,14 @@ public class ArrowStreamWriter implements AutoCloseable {
3535
* Creates the stream writer. non-blocking.
3636
* totalBatches can be set if the writer knows beforehand. Can be -1 if unknown.
3737
*/
38-
public ArrowStreamWriter(WritableByteChannel out, Schema schema, int totalBatches) {
38+
public ArrowStreamWriter(WritableByteChannel out, Schema schema) {
3939
this.out = new WriteChannel(out);
4040
this.schema = schema;
4141
}
4242

43-
public ArrowStreamWriter(OutputStream out, Schema schema, int totalBatches)
43+
public ArrowStreamWriter(OutputStream out, Schema schema)
4444
throws IOException {
45-
this(Channels.newChannel(out), schema, totalBatches);
45+
this(Channels.newChannel(out), schema);
4646
}
4747

4848
public long bytesWritten() { return out.getCurrentPosition(); }
@@ -53,6 +53,14 @@ public void writeRecordBatch(ArrowRecordBatch batch) throws IOException {
5353
MessageSerializer.serialize(out, batch);
5454
}
5555

56+
/**
57+
* End the stream. This is not required and this object can simply be closed.
58+
*/
59+
public void end() throws IOException {
60+
checkAndSendHeader();
61+
out.writeIntLittleEndian(0);
62+
}
63+
5664
@Override
5765
public void close() throws IOException {
5866
// The header might not have been sent if this is an empty stream. Send it even in

‎java/vector/src/main/java/org/apache/arrow/vector/stream/MessageSerializer.java‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -235,11 +235,10 @@ private static ByteBuffer serializeMessage(FlatBufferBuilder builder, byte heade
235235
private static Message deserializeMessage(ReadChannel in, byte headerType) throws IOException {
236236
// Read the message size. There is an i32 little endian prefix.
237237
ByteBuffer buffer = ByteBuffer.allocate(4);
238-
if (in.readFully(buffer) != 4) {
239-
return null;
240-
}
241-
238+
if (in.readFully(buffer) != 4) return null;
242239
int messageLength = bytesToInt(buffer.array());
240+
if (messageLength == 0) return null;
241+
243242
buffer = ByteBuffer.allocate(messageLength);
244243
if (in.readFully(buffer) != messageLength) {
245244
throw new IOException(

‎java/vector/src/test/java/org/apache/arrow/vector/file/TestArrowFile.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -232,7 +232,7 @@ public void testWriteReadMultipleRBs() throws IOException {
232232
Schema schema = vectorUnloader0.getSchema();
233233
Assert.assertEquals(2, schema.getFields().size());
234234
try (ArrowWriter arrowWriter = new ArrowWriter(fileOutputStream.getChannel(), schema);
235-
ArrowStreamWriter streamWriter = new ArrowStreamWriter(stream, schema, 2)) {
235+
ArrowStreamWriter streamWriter = new ArrowStreamWriter(stream, schema)) {
236236
try (ArrowRecordBatch recordBatch = vectorUnloader0.getRecordBatch()) {
237237
Assert.assertEquals("RB #0", counts[0], recordBatch.getLength());
238238
arrowWriter.writeRecordBatch(recordBatch);
@@ -399,7 +399,7 @@ private void write(FieldVector parent, File file, OutputStream outStream) throws
399399
// Also try serializing to the stream writer.
400400
if (outStream != null) {
401401
try (
402-
ArrowStreamWriter arrowWriter = new ArrowStreamWriter(outStream, schema, -1);
402+
ArrowStreamWriter arrowWriter = new ArrowStreamWriter(outStream, schema);
403403
ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch();
404404
) {
405405
arrowWriter.writeRecordBatch(recordBatch);

‎java/vector/src/test/java/org/apache/arrow/vector/stream/TestArrowStream.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ public void testEmptyStream() throws IOException {
4242

4343
// Write the stream.
4444
ByteArrayOutputStream out = new ByteArrayOutputStream();
45-
try (ArrowStreamWriter writer = new ArrowStreamWriter(out, schema, -1)) {
45+
try (ArrowStreamWriter writer = new ArrowStreamWriter(out, schema)) {
4646
}
4747

4848
ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());
@@ -66,7 +66,7 @@ public void testReadWrite() throws IOException {
6666
BufferAllocator alloc = new RootAllocator(Long.MAX_VALUE);
6767
ByteArrayOutputStream out = new ByteArrayOutputStream();
6868
long bytesWritten = 0;
69-
try (ArrowStreamWriter writer = new ArrowStreamWriter(out, schema, numBatches)) {
69+
try (ArrowStreamWriter writer = new ArrowStreamWriter(out, schema)) {
7070
ArrowBuf validityb = MessageSerializerTest.buf(alloc, validity);
7171
ArrowBuf valuesb = MessageSerializerTest.buf(alloc, values);
7272
for (int i = 0; i < numBatches; i++) {

‎java/vector/src/test/java/org/apache/arrow/vector/stream/TestArrowStreamPipe.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ private final class WriterThread extends Thread {
4747
public WriterThread(int numBatches, WritableByteChannel sinkChannel)
4848
throws IOException {
4949
this.numBatches = numBatches;
50-
writer = new ArrowStreamWriter(sinkChannel, schema, -1);
50+
writer = new ArrowStreamWriter(sinkChannel, schema);
5151
}
5252

5353
@Override

0 commit comments

Comments
 (0)