Skip to content

Commit fc58e3b

Browse files
sarahchen6sarahchen6
committed
Fallback to JNR
1 parent e71516e commit fc58e3b

2 files changed

Lines changed: 131 additions & 56 deletions

File tree

‎src/main/java/com/timgroup/statsd/UnixStreamClientChannel.java‎

Lines changed: 84 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -107,8 +107,11 @@ public int writeAll(ByteBuffer bb, boolean canReturnOnTimeout, long deadline)
107107
}
108108

109109
if (read == 0) {
110-
if (canReturnOnTimeout && written == 0 && delegate.isBlocking()) {
111-
return written;
110+
if (delegate.isBlocking()) {
111+
if (canReturnOnTimeout && written == 0) {
112+
return written;
113+
}
114+
throw new IOException("Write timed out");
112115
}
113116

114117
try (Selector selector = Selector.open()) {
@@ -159,60 +162,7 @@ private void connect() throws IOException {
159162
}
160163

161164
// Use native JDK support for UDS on Java 16+ and jnr-unixsocket otherwise
162-
if (VersionUtils.isJavaVersionAtLeast(16) && enableJdkSocket) {
163-
// Only SocketChannel.open(ProtocolFamily) needs reflection; connect/finishConnect
164-
// have existed since Java 1.4 and are called directly once we have the channel.
165-
SocketChannel channel = VersionUtils.openUnixSocketChannel();
166-
167-
SocketAddress connectAddress =
168-
address instanceof UnixSocketAddressWithTransport
169-
? ((UnixSocketAddressWithTransport) address).getAddress()
170-
: address;
171-
172-
try {
173-
if (bufferSize > 0) {
174-
channel.setOption(StandardSocketOptions.SO_SNDBUF, bufferSize);
175-
}
176-
if (connectionTimeout <= 0) {
177-
channel.configureBlocking(true);
178-
channel.connect(connectAddress);
179-
channel.configureBlocking(false);
180-
} else {
181-
channel.configureBlocking(false);
182-
long deadline = System.nanoTime() + connectionTimeout * 1_000_000L;
183-
if (channel.connect(connectAddress)) {
184-
this.delegate = channel;
185-
return;
186-
}
187-
188-
try (Selector selector = Selector.open()) {
189-
channel.register(selector, SelectionKey.OP_CONNECT);
190-
while (!channel.finishConnect()) {
191-
long remainingNs = deadline - System.nanoTime();
192-
if (remainingNs <= 0) {
193-
throw new IOException(
194-
"Connection timed out after " + connectionTimeout + "ms");
195-
}
196-
long selectTimeout = Math.max(1L, remainingNs / 1_000_000L);
197-
if (selector.select(selectTimeout) == 0
198-
&& System.nanoTime() >= deadline) {
199-
throw new IOException(
200-
"Connection timed out after " + connectionTimeout + "ms");
201-
}
202-
selector.selectedKeys().clear();
203-
}
204-
}
205-
}
206-
} catch (IOException e) {
207-
try {
208-
channel.close();
209-
} catch (IOException __) {
210-
// ignore
211-
}
212-
throw e;
213-
}
214-
215-
this.delegate = channel;
165+
if (VersionUtils.isJavaVersionAtLeast(16) && enableJdkSocket && connectWithJdkSocket()) {
216166
return;
217167
}
218168
// Default to jnr-unixsocket if Java version is < 16 or native support is disabled
@@ -258,6 +208,84 @@ private void connect() throws IOException {
258208
this.delegate = channel;
259209
}
260210

211+
private boolean connectWithJdkSocket() throws IOException {
212+
SocketChannel channel = null;
213+
SocketAddress connectAddress;
214+
215+
try {
216+
// Only SocketChannel.open(ProtocolFamily) needs reflection; connect/finishConnect
217+
// have existed since Java 1.4 and are called directly once we have the channel.
218+
channel = openJdkSocketChannel();
219+
connectAddress = nativeSocketAddress(address);
220+
if (bufferSize > 0) {
221+
channel.setOption(StandardSocketOptions.SO_SNDBUF, bufferSize);
222+
}
223+
} catch (Exception | LinkageError e) {
224+
closeQuietly(channel);
225+
return false;
226+
}
227+
228+
try {
229+
if (connectionTimeout <= 0) {
230+
channel.configureBlocking(true);
231+
channel.connect(connectAddress);
232+
channel.configureBlocking(false);
233+
} else {
234+
channel.configureBlocking(false);
235+
long deadline = System.nanoTime() + connectionTimeout * 1_000_000L;
236+
if (!channel.connect(connectAddress)) {
237+
try (Selector selector = Selector.open()) {
238+
channel.register(selector, SelectionKey.OP_CONNECT);
239+
while (!channel.finishConnect()) {
240+
long remainingNs = deadline - System.nanoTime();
241+
if (remainingNs <= 0) {
242+
throw new IOException(
243+
"Connection timed out after " + connectionTimeout + "ms");
244+
}
245+
long selectTimeout = Math.max(1L, remainingNs / 1_000_000L);
246+
if (selector.select(selectTimeout) == 0
247+
&& System.nanoTime() >= deadline) {
248+
throw new IOException(
249+
"Connection timed out after " + connectionTimeout + "ms");
250+
}
251+
selector.selectedKeys().clear();
252+
}
253+
}
254+
}
255+
}
256+
} catch (IOException | RuntimeException e) {
257+
closeQuietly(channel);
258+
throw e;
259+
}
260+
261+
this.delegate = channel;
262+
return true;
263+
}
264+
265+
SocketChannel openJdkSocketChannel() throws IOException {
266+
return VersionUtils.openUnixSocketChannel();
267+
}
268+
269+
private static SocketAddress nativeSocketAddress(SocketAddress address) {
270+
if (address instanceof UnixSocketAddressWithTransport) {
271+
address = ((UnixSocketAddressWithTransport) address).getAddress();
272+
}
273+
if (address instanceof UnixSocketAddress) {
274+
return VersionUtils.newUnixDomainSocketAddress(((UnixSocketAddress) address).path());
275+
}
276+
return address;
277+
}
278+
279+
private static void closeQuietly(SocketChannel channel) {
280+
if (channel != null) {
281+
try {
282+
channel.close();
283+
} catch (IOException ignored) {
284+
// ignore
285+
}
286+
}
287+
}
288+
261289
@Override
262290
public void close() throws IOException {
263291
disconnect();

‎src/test/java/com/timgroup/statsd/UnixStreamSocketTest.java‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010

1111
import java.io.File;
1212
import java.io.IOException;
13+
import java.nio.ByteBuffer;
14+
import java.nio.channels.SocketChannel;
1315
import java.nio.file.Files;
1416
import java.util.logging.Logger;
1517
import org.junit.After;
@@ -185,6 +187,39 @@ public void resist_dsd_timeout() throws Exception {
185187
server.clear();
186188
}
187189

190+
@Test(timeout = 5000L)
191+
public void nativeSocketAcceptsLegacyUnixAddress() throws Exception {
192+
Assume.assumeTrue(VersionUtils.isJavaVersionAtLeast(16));
193+
UnixStreamClientChannel channel =
194+
new UnixStreamClientChannel(
195+
new jnr.unixsocket.UnixSocketAddress(socketFile.getPath()),
196+
500,
197+
500,
198+
-1,
199+
true);
200+
201+
assertChannelSends(channel, "legacy.address:1|c");
202+
}
203+
204+
@Test(timeout = 5000L)
205+
public void fallsBackWhenNativeSocketCannotBeOpened() throws Exception {
206+
Assume.assumeTrue(VersionUtils.isJavaVersionAtLeast(16));
207+
UnixStreamClientChannel channel =
208+
new UnixStreamClientChannel(
209+
new jnr.unixsocket.UnixSocketAddress(socketFile.getPath()),
210+
500,
211+
500,
212+
-1,
213+
true) {
214+
@Override
215+
SocketChannel openJdkSocketChannel() throws IOException {
216+
throw new IOException("Native UDS unavailable");
217+
}
218+
};
219+
220+
assertChannelSends(channel, "fallback:1|c");
221+
}
222+
188223
@Test(timeout = 5000L)
189224
public void stream_uds_has_uds_buffer_size() throws Exception {
190225
final NonBlockingStatsDClient client =
@@ -211,4 +246,16 @@ public void max_packet_size_override() throws Exception {
211246

212247
assertEquals(client.statsDProcessor.bufferPool.getBufferSize(), 576);
213248
}
249+
250+
private static void assertChannelSends(UnixStreamClientChannel channel, String message)
251+
throws Exception {
252+
try {
253+
channel.write(ByteBuffer.wrap(message.getBytes(NonBlockingStatsDClient.UTF_8)));
254+
server.waitForMessage();
255+
assertThat(server.messagesReceived(), hasItem(message));
256+
server.clear();
257+
} finally {
258+
channel.close();
259+
}
260+
}
214261
}

0 commit comments

Comments
 (0)