Skip to content

Commit b87017c

Browse files
Tim-BrooksTim-Brooks
committed
Add test for CopyBytesSocketChannel (elastic#45873)
Currently we use a custom CopyBytesSocketChannel for interfacing with netty. We have integration tests that use this channel, however we never verify the read and write behavior in the face of potential partial writes. This commit adds a test for this behavior.
1 parent dd6c13f commit b87017c

File tree

2 files changed

+198
-3
lines changed

2 files changed

+198
-3
lines changed

‎modules/transport-netty4/src/main/java/org/elasticsearch/transport/CopyBytesSocketChannel.java‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
import io.netty.channel.socket.nio.NioSocketChannel;
4141
import org.elasticsearch.common.SuppressForbidden;
4242

43+
import java.io.IOException;
4344
import java.nio.ByteBuffer;
4445
import java.nio.channels.SocketChannel;
4546

@@ -74,7 +75,6 @@ public CopyBytesSocketChannel() {
7475

7576
@Override
7677
protected void doWrite(ChannelOutboundBuffer in) throws Exception {
77-
SocketChannel ch = javaChannel();
7878
int writeSpinCount = config().getWriteSpinCount();
7979
do {
8080
if (in.isEmpty()) {
@@ -99,7 +99,7 @@ protected void doWrite(ChannelOutboundBuffer in) throws Exception {
9999
ioBuffer.flip();
100100

101101
int attemptedBytes = ioBuffer.remaining();
102-
final int localWrittenBytes = ch.write(ioBuffer);
102+
final int localWrittenBytes = writeToSocketChannel(javaChannel(), ioBuffer);
103103
if (localWrittenBytes <= 0) {
104104
incompleteWrite(true);
105105
return;
@@ -119,14 +119,24 @@ protected int doReadBytes(ByteBuf byteBuf) throws Exception {
119119
final RecvByteBufAllocator.Handle allocHandle = unsafe().recvBufAllocHandle();
120120
allocHandle.attemptedBytesRead(byteBuf.writableBytes());
121121
ByteBuffer ioBuffer = getIoBuffer();
122-
int bytesRead = javaChannel().read(ioBuffer);
122+
int bytesRead = readFromSocketChannel(javaChannel(), ioBuffer);
123123
ioBuffer.flip();
124124
if (bytesRead > 0) {
125125
byteBuf.writeBytes(ioBuffer);
126126
}
127127
return bytesRead;
128128
}
129129

130+
// Protected so that tests can verify behavior and simulate partial writes
131+
protected int writeToSocketChannel(SocketChannel socketChannel, ByteBuffer ioBuffer) throws IOException {
132+
return socketChannel.write(ioBuffer);
133+
}
134+
135+
// Protected so that tests can verify behavior
136+
protected int readFromSocketChannel(SocketChannel socketChannel, ByteBuffer ioBuffer) throws IOException {
137+
return socketChannel.read(ioBuffer);
138+
}
139+
130140
private static ByteBuffer getIoBuffer() {
131141
ByteBuffer ioBuffer = CopyBytesSocketChannel.ioBuffer.get();
132142
ioBuffer.clear();

‎modules/transport-netty4/src/test/java/org/elasticsearch/transport/CopyBytesSocketChannelTests.java‎

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
/*
2+
* Licensed to Elasticsearch under one or more contributor
3+
* license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright
5+
* ownership. Elasticsearch licenses this file to you under
6+
* the Apache License, Version 2.0 (the "License"); you may
7+
* not use this file except in compliance with the License.
8+
* 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,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.elasticsearch.transport;
20+
21+
import io.netty.bootstrap.Bootstrap;
22+
import io.netty.bootstrap.ServerBootstrap;
23+
import io.netty.buffer.ByteBuf;
24+
import io.netty.buffer.Unpooled;
25+
import io.netty.buffer.UnpooledByteBufAllocator;
26+
import io.netty.channel.Channel;
27+
import io.netty.channel.ChannelFuture;
28+
import io.netty.channel.ChannelHandlerContext;
29+
import io.netty.channel.ChannelInitializer;
30+
import io.netty.channel.ChannelOption;
31+
import io.netty.channel.SimpleChannelInboundHandler;
32+
import io.netty.channel.nio.NioEventLoopGroup;
33+
import org.elasticsearch.common.SuppressForbidden;
34+
import org.elasticsearch.test.ESTestCase;
35+
36+
import java.io.IOException;
37+
import java.net.InetAddress;
38+
import java.net.InetSocketAddress;
39+
import java.nio.ByteBuffer;
40+
import java.nio.channels.SocketChannel;
41+
import java.nio.charset.StandardCharsets;
42+
import java.util.concurrent.ConcurrentLinkedQueue;
43+
import java.util.concurrent.TimeUnit;
44+
import java.util.concurrent.atomic.AtomicInteger;
45+
import java.util.concurrent.atomic.AtomicReference;
46+
47+
public class CopyBytesSocketChannelTests extends ESTestCase {
48+
49+
private final UnpooledByteBufAllocator alloc = new UnpooledByteBufAllocator(false);
50+
private final AtomicReference<CopyBytesSocketChannel> accepted = new AtomicReference<>();
51+
private final AtomicInteger serverBytesReceived = new AtomicInteger();
52+
private final AtomicInteger clientBytesReceived = new AtomicInteger();
53+
private final ConcurrentLinkedQueue<ByteBuf> serverReceived = new ConcurrentLinkedQueue<>();
54+
private final ConcurrentLinkedQueue<ByteBuf> clientReceived = new ConcurrentLinkedQueue<>();
55+
private NioEventLoopGroup eventLoopGroup;
56+
private InetSocketAddress serverAddress;
57+
private Channel serverChannel;
58+
59+
@Override
60+
@SuppressForbidden(reason = "calls getLocalHost")
61+
public void setUp() throws Exception {
62+
super.setUp();
63+
eventLoopGroup = new NioEventLoopGroup(1);
64+
ServerBootstrap serverBootstrap = new ServerBootstrap();
65+
serverBootstrap.channel(CopyBytesServerSocketChannel.class);
66+
serverBootstrap.group(eventLoopGroup);
67+
serverBootstrap.option(ChannelOption.ALLOCATOR, alloc);
68+
serverBootstrap.childOption(ChannelOption.ALLOCATOR, alloc);
69+
serverBootstrap.childHandler(new ChannelInitializer<>() {
70+
@Override
71+
protected void initChannel(Channel ch) {
72+
accepted.set((CopyBytesSocketChannel) ch);
73+
ch.pipeline().addLast(new SimpleChannelInboundHandler<>() {
74+
@Override
75+
protected void channelRead0(ChannelHandlerContext ctx, Object msg) {
76+
ByteBuf buffer = (ByteBuf) msg;
77+
serverBytesReceived.addAndGet(buffer.readableBytes());
78+
serverReceived.add(buffer.retain());
79+
}
80+
});
81+
}
82+
});
83+
84+
ChannelFuture bindFuture = serverBootstrap.bind(new InetSocketAddress(InetAddress.getLocalHost(), 0));
85+
assertTrue(bindFuture.await(10, TimeUnit.SECONDS));
86+
serverAddress = (InetSocketAddress) bindFuture.channel().localAddress();
87+
bindFuture.isSuccess();
88+
serverChannel = bindFuture.channel();
89+
}
90+
91+
@Override
92+
public void tearDown() throws Exception {
93+
super.tearDown();
94+
try {
95+
assertTrue(serverChannel.close().await(10, TimeUnit.SECONDS));
96+
} finally {
97+
eventLoopGroup.shutdownGracefully().await(10, TimeUnit.SECONDS);
98+
}
99+
}
100+
101+
public void testSendAndReceive() throws Exception {
102+
final Bootstrap bootstrap = new Bootstrap();
103+
bootstrap.group(eventLoopGroup);
104+
bootstrap.channel(VerifyingCopyChannel.class);
105+
bootstrap.option(ChannelOption.ALLOCATOR, alloc);
106+
bootstrap.handler(new ChannelInitializer<>() {
107+
@Override
108+
protected void initChannel(Channel ch) {
109+
ch.pipeline().addLast(new SimpleChannelInboundHandler<>() {
110+
@Override
111+
protected void channelRead0(ChannelHandlerContext ctx, Object msg) {
112+
ByteBuf buffer = (ByteBuf) msg;
113+
clientBytesReceived.addAndGet(buffer.readableBytes());
114+
clientReceived.add(buffer.retain());
115+
}
116+
});
117+
}
118+
});
119+
120+
ChannelFuture connectFuture = bootstrap.connect(serverAddress);
121+
connectFuture.await(10, TimeUnit.SECONDS);
122+
assertTrue(connectFuture.isSuccess());
123+
CopyBytesSocketChannel copyChannel = (CopyBytesSocketChannel) connectFuture.channel();
124+
ByteBuf clientData = generateData();
125+
ByteBuf serverData = generateData();
126+
127+
try {
128+
assertBusy(() -> assertNotNull(accepted.get()));
129+
int clientBytesToWrite = clientData.readableBytes();
130+
ChannelFuture clientWriteFuture = copyChannel.writeAndFlush(clientData.retainedSlice());
131+
clientWriteFuture.await(10, TimeUnit.SECONDS);
132+
assertBusy(() -> assertEquals(clientBytesToWrite, serverBytesReceived.get()));
133+
134+
int serverBytesToWrite = serverData.readableBytes();
135+
ChannelFuture serverWriteFuture = accepted.get().writeAndFlush(serverData.retainedSlice());
136+
assertTrue(serverWriteFuture.await(10, TimeUnit.SECONDS));
137+
assertBusy(() -> assertEquals(serverBytesToWrite, clientBytesReceived.get()));
138+
139+
ByteBuf compositeServerReceived = Unpooled.wrappedBuffer(serverReceived.toArray(new ByteBuf[0]));
140+
assertEquals(clientData, compositeServerReceived);
141+
ByteBuf compositeClientReceived = Unpooled.wrappedBuffer(clientReceived.toArray(new ByteBuf[0]));
142+
assertEquals(serverData, compositeClientReceived);
143+
} finally {
144+
clientData.release();
145+
serverData.release();
146+
serverReceived.forEach(ByteBuf::release);
147+
clientReceived.forEach(ByteBuf::release);
148+
assertTrue(copyChannel.close().await(10, TimeUnit.SECONDS));
149+
}
150+
}
151+
152+
private ByteBuf generateData() {
153+
return Unpooled.wrappedBuffer(randomAlphaOfLength(randomIntBetween(1 << 22, 1 << 23)).getBytes(StandardCharsets.UTF_8));
154+
}
155+
156+
public static class VerifyingCopyChannel extends CopyBytesSocketChannel {
157+
158+
public VerifyingCopyChannel() {
159+
super();
160+
}
161+
162+
@Override
163+
protected int writeToSocketChannel(SocketChannel socketChannel, ByteBuffer ioBuffer) throws IOException {
164+
assertTrue("IO Buffer must be a direct byte buffer", ioBuffer.isDirect());
165+
int remaining = ioBuffer.remaining();
166+
int originalLimit = ioBuffer.limit();
167+
// If greater than a KB, possibly invoke a partial write.
168+
if (remaining > 1024) {
169+
if (randomBoolean()) {
170+
int bytes = randomIntBetween(remaining / 2, remaining);
171+
ioBuffer.limit(ioBuffer.position() + bytes);
172+
}
173+
}
174+
int written = socketChannel.write(ioBuffer);
175+
ioBuffer.limit(originalLimit);
176+
return written;
177+
}
178+
179+
@Override
180+
protected int readFromSocketChannel(SocketChannel socketChannel, ByteBuffer ioBuffer) throws IOException {
181+
assertTrue("IO Buffer must be a direct byte buffer", ioBuffer.isDirect());
182+
return socketChannel.read(ioBuffer);
183+
}
184+
}
185+
}

0 commit comments

Comments
 (0)