diff --git a/dd-trace-core/build.gradle b/dd-trace-core/build.gradle index 5424650d961..8b762e2461e 100644 --- a/dd-trace-core/build.gradle +++ b/dd-trace-core/build.gradle @@ -116,6 +116,8 @@ dependencies { testImplementation group: 'commons-codec', name: 'commons-codec', version: '1.3' testImplementation group: 'com.amazonaws', name: 'aws-lambda-java-events', version:'3.11.0' testImplementation group: 'com.google.protobuf', name: 'protobuf-java', version: '3.14.0' + testImplementation libs.jnr.unixsocket + testImplementation group: 'com.squareup.okhttp3', name: 'mockwebserver', version: libs.versions.okhttp.legacy.get() testImplementation libs.testcontainers testImplementation project(':utils:test-junit-utils') testImplementation project(':utils:test-junit-converter-utils') diff --git a/dd-trace-core/src/test/java/datadog/common/socket/UnixDomainServerSocketFactory.java b/dd-trace-core/src/test/java/datadog/common/socket/UnixDomainServerSocketFactory.java new file mode 100644 index 00000000000..6cad92ba460 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/common/socket/UnixDomainServerSocketFactory.java @@ -0,0 +1,109 @@ +package datadog.common.socket; + +import java.io.File; +import java.io.IOException; +import java.net.InetAddress; +import java.net.InetSocketAddress; +import java.net.ServerSocket; +import java.net.Socket; +import java.net.SocketAddress; +import java.net.SocketException; +import java.nio.channels.ClosedChannelException; +import javax.net.ServerSocketFactory; +import jnr.unixsocket.UnixServerSocketChannel; +import jnr.unixsocket.UnixSocketAddress; +import jnr.unixsocket.UnixSocketChannel; + +/** + * Adapts a JNR Unix-domain server channel to APIs such as MockWebServer that require a {@link + * ServerSocket}. Adapted from OkHttp's Unix-domain + * socket sample. + */ +public final class UnixDomainServerSocketFactory extends ServerSocketFactory { + private final File path; + + public UnixDomainServerSocketFactory(File path) { + this.path = path; + } + + @Override + public ServerSocket createServerSocket() throws IOException { + return new UnixDomainServerSocket(); + } + + @Override + public ServerSocket createServerSocket(int port) throws IOException { + return createServerSocket(); + } + + @Override + public ServerSocket createServerSocket(int port, int backlog) throws IOException { + return createServerSocket(); + } + + @Override + public ServerSocket createServerSocket(int port, int backlog, InetAddress inetAddress) + throws IOException { + return createServerSocket(); + } + + private final class UnixDomainServerSocket extends ServerSocket { + private UnixServerSocketChannel serverSocketChannel; + private InetSocketAddress endpoint; + + private UnixDomainServerSocket() throws IOException {} + + @Override + public void bind(SocketAddress endpoint, int backlog) throws IOException { + this.endpoint = (InetSocketAddress) endpoint; + UnixServerSocketChannel channel = UnixServerSocketChannel.open(); + boolean bound = false; + try { + channel.configureBlocking(true); + channel.socket().bind(new UnixSocketAddress(path)); + serverSocketChannel = channel; + bound = true; + } finally { + if (!bound) { + channel.close(); + } + } + } + + @Override + public void setReuseAddress(boolean on) { + // MockWebServer configures this TCP option before binding. It has no UDS equivalent. + } + + @Override + public int getLocalPort() { + return 1; // MockWebServer requires a port even though a UDS has none. + } + + @Override + public SocketAddress getLocalSocketAddress() { + return endpoint; + } + + @Override + public Socket accept() throws IOException { + try { + UnixSocketChannel channel = serverSocketChannel.accept(); + return new TunnelingUnixSocket(path, channel, endpoint); + } catch (ClosedChannelException e) { + SocketException socketException = new SocketException("Socket is closed"); + socketException.initCause(e); + throw socketException; + } + } + + @Override + public void close() throws IOException { + super.close(); + if (serverSocketChannel != null) { + serverSocketChannel.close(); + } + } + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/common/writer/DDAgentWriterCombinedTest.java b/dd-trace-core/src/test/java/datadog/trace/common/writer/DDAgentWriterCombinedTest.java index ae9741fe827..b0d970a7960 100644 --- a/dd-trace-core/src/test/java/datadog/trace/common/writer/DDAgentWriterCombinedTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/common/writer/DDAgentWriterCombinedTest.java @@ -2,9 +2,18 @@ import static datadog.trace.api.ProtocolVersion.V0_5; import static datadog.trace.api.config.GeneralConfig.EXPERIMENTAL_PROPAGATE_PROCESS_TAGS_ENABLED; +import static datadog.trace.api.config.GeneralConfig.JDK_SOCKET_ENABLED; import static datadog.trace.common.writer.ddagent.Prioritization.ENSURE_TRACE; +import static okhttp3.mockwebserver.SocketPolicy.DISCONNECT_AT_END; +import static okhttp3.mockwebserver.SocketPolicy.NO_RESPONSE; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assumptions.assumeTrue; +import static org.junit.jupiter.api.condition.JRE.JAVA_16; +import static org.junit.jupiter.api.condition.OS.LINUX; +import static org.junit.jupiter.api.condition.OS.MAC; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyLong; @@ -18,6 +27,7 @@ import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; +import datadog.common.socket.UnixDomainServerSocketFactory; import datadog.communication.ddagent.DDAgentFeaturesDiscovery; import datadog.communication.http.OkHttpUtils; import datadog.communication.serialization.FlushingBuffer; @@ -39,6 +49,8 @@ import datadog.trace.test.junit.utils.config.WithConfig; import datadog.trace.test.util.Flaky; import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; import java.util.Collections; import java.util.List; import java.util.concurrent.CountDownLatch; @@ -46,11 +58,17 @@ import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import okhttp3.HttpUrl; +import okhttp3.mockwebserver.MockResponse; +import okhttp3.mockwebserver.MockWebServer; +import okhttp3.mockwebserver.RecordedRequest; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.condition.EnabledForJreRange; +import org.junit.jupiter.api.condition.EnabledOnOs; import org.mockito.Mockito; import org.tabletest.junit.TableTest; @@ -336,12 +354,11 @@ void monitorHappyPath(String agentVersion) { List minimalTrace = createMinimalTrace(); // DQH -- need to set-up a dummy agent for the final send callback to work - JavaTestHttpServer agent = + try (JavaTestHttpServer agent = JavaTestHttpServer.httpServer( server -> server.handlers( - h -> h.put(agentVersion, api -> api.getResponse().status(200).send()))); - try { + h -> h.put(agentVersion, api -> api.getResponse().status(200).send())))) { HttpUrl agentUrl = HttpUrl.get(agent.getAddress()); okhttp3.OkHttpClient client = OkHttpUtils.buildHttpClient(agentUrl, 1000); DDAgentFeaturesDiscovery discovery = @@ -380,8 +397,6 @@ void monitorHappyPath(String agentVersion) { writer.close(); verify(healthMetrics, times(1)).onShutdown(true); - } finally { - agent.close(); } } @@ -397,7 +412,9 @@ void monitorAgentReturnsError(String agentVersion) { // DQH -- need to set-up a dummy agent for the final send callback to work final boolean[] first = {true}; - JavaTestHttpServer agent = + // DQH - DDApi sniffs for end point existence, so respond with 200 the + // first time + try (JavaTestHttpServer agent = JavaTestHttpServer.httpServer( server -> server.handlers( @@ -413,8 +430,7 @@ void monitorAgentReturnsError(String agentVersion) { } else { api.getResponse().status(500).send(); } - }))); - try { + })))) { HttpUrl agentUrl = HttpUrl.get(agent.getAddress()); okhttp3.OkHttpClient client = OkHttpUtils.buildHttpClient(agentUrl, 1000); DDAgentFeaturesDiscovery discovery = @@ -453,8 +469,82 @@ void monitorAgentReturnsError(String agentVersion) { writer.close(); verify(healthMetrics, times(1)).onShutdown(true); + } + } + + @Test + @WithConfig(key = JDK_SOCKET_ENABLED, value = "true") + @EnabledForJreRange(min = JAVA_16) + @EnabledOnOs({LINUX, MAC}) + void unixSocketTimeoutKeepsWorkerAliveAndReconnects() throws Exception { + assertTrue(Config.get().isJdkSocketEnabled()); + + Path socketPath = Files.createTempFile("dd-trace-agent-", ".sock"); + Files.delete(socketPath); + + HealthMetrics healthMetrics = mock(HealthMetrics.class); + AtomicReference failedSendThread = new AtomicReference<>(); + AtomicReference successfulSendThread = new AtomicReference<>(); + doAnswer( + invocation -> { + failedSendThread.set(Thread.currentThread()); + return null; + }) + .when(healthMetrics) + .onFailedSend(anyInt(), anyInt(), any()); + doAnswer( + invocation -> { + successfulSendThread.set(Thread.currentThread()); + return null; + }) + .when(healthMetrics) + .onSend(anyInt(), anyInt(), any()); + + DDAgentFeaturesDiscovery discovery = mock(DDAgentFeaturesDiscovery.class); + when(discovery.getTraceEndpoint()).thenReturn("v0.4/traces"); + + try (MockWebServer server = new MockWebServer(); + DDAgentWriter writer = + DDAgentWriter.builder() + .featureDiscovery(discovery) + .unixDomainSocket(socketPath.toString()) + .timeoutMillis(100) + .monitoring(monitoring) + .healthMetrics(healthMetrics) + .flushIntervalMilliseconds(-1) + .flushTimeout(5, TimeUnit.SECONDS) + .build()) { + server.setServerSocketFactory(new UnixDomainServerSocketFactory(socketPath.toFile())); + // Read the first request fully, then withhold the response to trigger a header-read timeout. + server.enqueue(new MockResponse().setSocketPolicy(NO_RESPONSE)); + server.enqueue(new MockResponse().setResponseCode(200).setSocketPolicy(DISCONNECT_AT_END)); + server.start(); + + writer.start(); + + writer.write(createMinimalTrace()); + assertTrue(writer.flush()); + + RecordedRequest failedRequest = server.takeRequest(5, TimeUnit.SECONDS); + assertNotNull(failedRequest); + assertEquals(0, failedRequest.getSequenceNumber()); + assertEquals(1, server.getRequestCount()); + verify(healthMetrics, times(1)).onFailedSend(anyInt(), anyInt(), any()); + + writer.write(createMinimalTrace()); + assertTrue(writer.flush()); + + RecordedRequest successfulRequest = server.takeRequest(5, TimeUnit.SECONDS); + assertNotNull(successfulRequest); + // Sequence numbers are per connection; zero again proves that this used a fresh socket. + assertEquals(0, successfulRequest.getSequenceNumber()); + assertEquals(2, server.getRequestCount()); + verify(healthMetrics, times(1)).onSend(anyInt(), anyInt(), any()); + assertSame(failedSendThread.get(), successfulSendThread.get()); + } catch (IOException | UnsupportedOperationException e) { + assumeTrue(false, "Unix-domain sockets are not supported: " + e.getMessage()); } finally { - agent.close(); + Files.deleteIfExists(socketPath); } } diff --git a/utils/socket-utils/build.gradle.kts b/utils/socket-utils/build.gradle.kts index 80055d2b62d..124ed80038a 100644 --- a/utils/socket-utils/build.gradle.kts +++ b/utils/socket-utils/build.gradle.kts @@ -12,6 +12,7 @@ extensions.getByName("tracerJava").withGroovyBuilder { } dependencies { + add("main_java17CompileOnly", project(":components:annotations")) implementation(project(":components:environment")) implementation(project(":utils:logging-utils")) implementation(libs.slf4j) diff --git a/utils/socket-utils/src/main/java/datadog/common/socket/UnixDomainSocketFactory.java b/utils/socket-utils/src/main/java/datadog/common/socket/UnixDomainSocketFactory.java index 1df84896c56..09d03d1325d 100644 --- a/utils/socket-utils/src/main/java/datadog/common/socket/UnixDomainSocketFactory.java +++ b/utils/socket-utils/src/main/java/datadog/common/socket/UnixDomainSocketFactory.java @@ -45,7 +45,7 @@ public Socket createSocket() throws IOException { if (this.useJdkUdsSocket) { try { return new TunnelingJdkSocket(this.path.toPath()); - } catch (Throwable ignore) { + } catch (IOException | UnsupportedOperationException ignore) { // fall back to jnr-unixsocket library } } diff --git a/utils/socket-utils/src/main/java17/datadog/common/socket/TunnelingJdkSocket.java b/utils/socket-utils/src/main/java17/datadog/common/socket/TunnelingJdkSocket.java index fb25228e0f6..26cafc49710 100644 --- a/utils/socket-utils/src/main/java17/datadog/common/socket/TunnelingJdkSocket.java +++ b/utils/socket-utils/src/main/java17/datadog/common/socket/TunnelingJdkSocket.java @@ -1,5 +1,6 @@ package datadog.common.socket; +import datadog.trace.api.internal.VisibleForTesting; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; @@ -8,8 +9,11 @@ import java.net.Socket; import java.net.SocketAddress; import java.net.SocketException; +import java.net.StandardProtocolFamily; import java.net.UnixDomainSocketAddress; import java.nio.ByteBuffer; +import java.nio.channels.CancelledKeyException; +import java.nio.channels.ClosedSelectorException; import java.nio.channels.SelectionKey; import java.nio.channels.Selector; import java.nio.channels.SocketChannel; @@ -25,34 +29,30 @@ * 16. */ final class TunnelingJdkSocket extends Socket { - private final SocketAddress unixSocketAddress; - private InetSocketAddress inetSocketAddress; + private final UnixDomainSocketAddress unixSocketAddress; + private final SocketChannel unixSocketChannel; - private SocketChannel unixSocketChannel; - private Selector selector; + private volatile InetSocketAddress inetSocketAddress; + @VisibleForTesting volatile Selector selector; - private int timeout; - private boolean shutIn; - private boolean shutOut; - private boolean closed; + private volatile int timeout; + private volatile boolean shutIn; + private volatile boolean shutOut; + private volatile boolean closed; static final int DEFAULT_BUFFER_SIZE = 8192; // Indicate that the buffer size is not set by initializing to -1 private int sendBufferSize = -1; private int receiveBufferSize = -1; - TunnelingJdkSocket(final Path path) { + TunnelingJdkSocket(final Path path) throws IOException, UnsupportedOperationException { this.unixSocketAddress = UnixDomainSocketAddress.of(path); - } - - TunnelingJdkSocket(final Path path, final InetSocketAddress address) { - this(path); - inetSocketAddress = address; + this.unixSocketChannel = SocketChannel.open(StandardProtocolFamily.UNIX); } @Override public boolean isConnected() { - return null != unixSocketChannel; + return inetSocketAddress != null; } @Override @@ -71,7 +71,7 @@ public boolean isClosed() { } @Override - public synchronized void setSoTimeout(int timeout) throws SocketException { + public void setSoTimeout(int timeout) throws SocketException { if (isClosed()) { throw new SocketException("Socket is closed"); } @@ -82,7 +82,7 @@ public synchronized void setSoTimeout(int timeout) throws SocketException { } @Override - public synchronized int getSoTimeout() throws SocketException { + public int getSoTimeout() throws SocketException { if (isClosed()) { throw new SocketException("Socket is closed"); } @@ -91,17 +91,7 @@ public synchronized int getSoTimeout() throws SocketException { @Override public void connect(final SocketAddress endpoint) throws IOException { - if (endpoint == null) { - throw new IllegalArgumentException("Endpoint cannot be null"); - } - if (isClosed()) { - throw new SocketException("Socket is closed"); - } - if (isConnected()) { - throw new SocketException("Socket is already connected"); - } - inetSocketAddress = (InetSocketAddress) endpoint; - unixSocketChannel = SocketChannel.open(unixSocketAddress); + connect(endpoint, 0); } // `timeout` is intentionally ignored here, like in the jnr-unixsocket implementation. @@ -121,8 +111,14 @@ public void connect(final SocketAddress endpoint, final int timeout) throws IOEx if (isConnected()) { throw new SocketException("Socket is already connected"); } - inetSocketAddress = (InetSocketAddress) endpoint; - unixSocketChannel = SocketChannel.open(unixSocketAddress); + InetSocketAddress inetSocketAddress = (InetSocketAddress) endpoint; + try { + unixSocketChannel.connect(unixSocketAddress); + this.inetSocketAddress = inetSocketAddress; + } catch (IOException e) { + close(); + throw e; + } } @Override @@ -200,69 +196,100 @@ public int getStreamBufferSize() throws SocketException { @Override public InputStream getInputStream() throws IOException { - if (isClosed()) { - throw new SocketException("Socket is closed"); - } - if (!isConnected()) { - throw new SocketException("Socket is not connected"); - } - if (isInputShutdown()) { - throw new SocketException("Socket input is shutdown"); - } - - if (selector == null) { - selector = Selector.open(); - unixSocketChannel.configureBlocking(false); - unixSocketChannel.register(selector, SelectionKey.OP_READ); - } - - return new InputStream() { - private final ByteBuffer buffer = ByteBuffer.allocate(getStreamBufferSize()); - - @Override - public int read() throws IOException { - byte[] nextByte = new byte[1]; - return (read(nextByte, 0, 1) == -1) ? -1 : (nextByte[0] & 0xFF); + // Serialize validation and selector publication with close() so close cannot miss a selector + // that is still being initialized. + synchronized (this) { + if (isClosed()) { + throw new SocketException("Socket is closed"); + } + if (!isConnected()) { + throw new SocketException("Socket is not connected"); + } + if (isInputShutdown()) { + throw new SocketException("Socket input is shutdown"); } - @Override - public int read(byte[] b, int off, int len) throws IOException { - if (isInputShutdown()) { - return -1; + Selector currentSelector = selector; + if (currentSelector == null) { + currentSelector = Selector.open(); + try { + unixSocketChannel.configureBlocking(false); + unixSocketChannel.register(currentSelector, SelectionKey.OP_READ); + selector = currentSelector; + } catch (IOException | RuntimeException e) { + try { + currentSelector.close(); + } catch (IOException closeException) { + e.addSuppressed(closeException); + } + throw e; } - buffer.clear(); + } + + return new InputStream() { + private final ByteBuffer buffer = ByteBuffer.allocate(getStreamBufferSize()); - int readyChannels = selector.select(timeout); - if (readyChannels == 0) { - return 0; + @Override + public int read() throws IOException { + byte[] nextByte = new byte[1]; + return (read(nextByte, 0, 1) == -1) ? -1 : (nextByte[0] & 0xFF); } - Set selectedKeys = selector.selectedKeys(); - synchronized (selectedKeys) { - Iterator keyIterator = selectedKeys.iterator(); - while (keyIterator.hasNext()) { - SelectionKey key = keyIterator.next(); - keyIterator.remove(); - if (key.isReadable()) { - int r = unixSocketChannel.read(buffer); - if (r == -1) { - return -1; + @Override + public int read(byte[] b, int off, int len) throws IOException { + if (isInputShutdown()) { + return -1; + } + buffer.clear(); + + Selector currentSelector = selector; + if (currentSelector == null) { + throw new SocketException("Socket is closed"); + } + + try { + int readyChannels = currentSelector.select(timeout); + if (readyChannels == 0) { + if (isClosed() || !currentSelector.isOpen()) { + throw new SocketException("Socket is closed"); } - buffer.flip(); - len = Math.min(r, len); - buffer.get(b, off, len); - return len; + return 0; } + + Set selectedKeys = currentSelector.selectedKeys(); + // Multiple input streams share this selector, so serialize iteration and removal from + // its non-thread-safe selected-key set. + synchronized (selectedKeys) { + Iterator keyIterator = selectedKeys.iterator(); + while (keyIterator.hasNext()) { + SelectionKey key = keyIterator.next(); + keyIterator.remove(); + if (key.isReadable()) { + int r = unixSocketChannel.read(buffer); + if (r == -1) { + return -1; + } + buffer.flip(); + len = Math.min(r, len); + buffer.get(b, off, len); + return len; + } + } + } + return 0; + } catch (ClosedSelectorException | CancelledKeyException e) { + SocketException socketException = new SocketException("Socket is closed"); + socketException.initCause(e); + throw socketException; } } - return 0; - } - @Override - public void close() throws IOException { - TunnelingJdkSocket.this.close(); - } - }; + @Override + public void close() throws IOException { + TunnelingJdkSocket.this.close(); + } + }; + } } @Override @@ -304,32 +331,38 @@ public void close() throws IOException { @Override public void shutdownInput() throws IOException { - if (isClosed()) { - throw new SocketException("Socket is closed"); - } - if (!isConnected()) { - throw new SocketException("Socket is not connected"); - } - if (isInputShutdown()) { - throw new SocketException("Socket input is already shutdown"); + // Keep validation, channel shutdown, and state publication atomic with close(). + synchronized (this) { + if (isClosed()) { + throw new SocketException("Socket is closed"); + } + if (!isConnected()) { + throw new SocketException("Socket is not connected"); + } + if (isInputShutdown()) { + throw new SocketException("Socket input is already shutdown"); + } + unixSocketChannel.shutdownInput(); + shutIn = true; } - unixSocketChannel.shutdownInput(); - shutIn = true; } @Override public void shutdownOutput() throws IOException { - if (isClosed()) { - throw new SocketException("Socket is closed"); - } - if (!isConnected()) { - throw new SocketException("Socket is not connected"); - } - if (isOutputShutdown()) { - throw new SocketException("Socket output is already shutdown"); + // Keep validation, channel shutdown, and state publication atomic with close(). + synchronized (this) { + if (isClosed()) { + throw new SocketException("Socket is closed"); + } + if (!isConnected()) { + throw new SocketException("Socket is not connected"); + } + if (isOutputShutdown()) { + throw new SocketException("Socket output is already shutdown"); + } + unixSocketChannel.shutdownOutput(); + shutOut = true; } - unixSocketChannel.shutdownOutput(); - shutOut = true; } @Override @@ -341,36 +374,29 @@ public InetAddress getInetAddress() { } @Override - public void close() throws IOException { - if (isClosed()) { - return; - } - // Ignore possible exceptions so that we continue closing the socket - try { - if (!isInputShutdown()) { - shutdownInput(); + public void close() { + Selector currentSelector; + // Publish the terminal state and snapshot the selector atomically with selector creation and + // half-close operations. The resources are closed after releasing this monitor. + synchronized (this) { + if (isClosed()) { + return; } - } catch (IOException e) { - } - try { - if (!isOutputShutdown()) { - shutdownOutput(); - } - } catch (IOException e) { + shutIn = true; + shutOut = true; + closed = true; + currentSelector = selector; } + // Ignore possible exceptions so that we continue closing the socket try { - if (selector != null) { - selector.close(); - selector = null; + if (currentSelector != null) { + currentSelector.close(); } - } catch (IOException e) { + } catch (IOException ignored) { } try { - if (unixSocketChannel != null) { - unixSocketChannel.close(); - } - } catch (IOException e) { + unixSocketChannel.close(); + } catch (IOException ignored) { } - closed = true; } } diff --git a/utils/socket-utils/src/test/java/datadog/common/socket/TunnelingJdkSocketTest.java b/utils/socket-utils/src/test/java/datadog/common/socket/TunnelingJdkSocketTest.java index 1c7c7ec7a19..6e55cdcd31d 100644 --- a/utils/socket-utils/src/test/java/datadog/common/socket/TunnelingJdkSocketTest.java +++ b/utils/socket-utils/src/test/java/datadog/common/socket/TunnelingJdkSocketTest.java @@ -2,6 +2,8 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -11,32 +13,60 @@ import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; -import java.lang.management.ManagementFactory; import java.net.InetSocketAddress; import java.net.SocketException; import java.net.StandardProtocolFamily; import java.net.UnixDomainSocketAddress; +import java.nio.channels.CancelledKeyException; +import java.nio.channels.ClosedSelectorException; +import java.nio.channels.SelectableChannel; +import java.nio.channels.SelectionKey; +import java.nio.channels.Selector; import java.nio.channels.ServerSocketChannel; import java.nio.channels.SocketChannel; +import java.nio.channels.spi.SelectorProvider; import java.nio.file.Files; import java.nio.file.Path; import java.time.Duration; -import java.util.concurrent.atomic.AtomicBoolean; +import java.util.Collections; +import java.util.HashSet; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledForJreRange; +@EnabledForJreRange(min = JAVA_16) public class TunnelingJdkSocketTest { - private static final AtomicBoolean isServerRunning = new AtomicBoolean(false); + private TestUnixSocketServer server; + + @BeforeAll + static void assumeUnixDomainSocketsAreSupported() { + Assumptions.assumeTrue(udsSupported()); + } + + @AfterEach + void closeServer() throws Exception { + if (server != null) { + server.close(); + server = null; + } + } @Test - @EnabledForJreRange(min = JAVA_16) public void testSocketConnectAndClose() throws Exception { Path socketPath = getSocketPath(); UnixDomainSocketAddress socketAddress = UnixDomainSocketAddress.of(socketPath); startServer(socketAddress); TunnelingJdkSocket clientSocket = new TunnelingJdkSocket(socketPath); + assertNotNull(clientSocket.getChannel()); + assertTrue(clientSocket.getChannel().isOpen()); assertFalse(clientSocket.isConnected()); assertFalse(clientSocket.isClosed()); @@ -55,6 +85,7 @@ public void testSocketConnectAndClose() throws Exception { assertTrue(clientSocket.isConnected()); assertTrue(clientSocket.isClosed()); + assertFalse(clientSocket.getChannel().isOpen()); assertTrue(clientSocket.isInputShutdown()); assertTrue(clientSocket.isOutputShutdown()); assertEquals(-1, inputStream.read()); @@ -62,12 +93,26 @@ public void testSocketConnectAndClose() throws Exception { assertThrows(SocketException.class, clientSocket::getInputStream); assertThrows(SocketException.class, clientSocket::getOutputStream); clientSocket.close(); + } + + @Test + public void testConnectFailureClosesSocket() throws Exception { + Path missingSocketPath = getSocketPath(); + TunnelingJdkSocket clientSocket = new TunnelingJdkSocket(missingSocketPath); + + assertThrows( + IOException.class, () -> clientSocket.connect(new InetSocketAddress("localhost", 0))); - isServerRunning.set(false); + assertFalse(clientSocket.isConnected()); + assertTrue(clientSocket.isClosed()); + assertTrue(clientSocket.isInputShutdown()); + assertTrue(clientSocket.isOutputShutdown()); + assertFalse(clientSocket.getChannel().isOpen()); + assertThrows( + SocketException.class, () -> clientSocket.connect(new InetSocketAddress("localhost", 0))); } @Test - @EnabledForJreRange(min = JAVA_16) public void testInputStreamClose() throws Exception { TunnelingJdkSocket clientSocket = createClient(); InputStream inputStream = clientSocket.getInputStream(); @@ -86,12 +131,9 @@ public void testInputStreamClose() throws Exception { assertThrows(IOException.class, () -> outputStream.write(1)); assertThrows(SocketException.class, clientSocket::getInputStream); assertThrows(SocketException.class, clientSocket::getOutputStream); - - isServerRunning.set(false); } @Test - @EnabledForJreRange(min = JAVA_16) public void testOutputStreamClose() throws Exception { TunnelingJdkSocket clientSocket = createClient(); InputStream inputStream = clientSocket.getInputStream(); @@ -110,12 +152,9 @@ public void testOutputStreamClose() throws Exception { assertThrows(IOException.class, () -> outputStream.write(1)); assertThrows(SocketException.class, clientSocket::getInputStream); assertThrows(SocketException.class, clientSocket::getOutputStream); - - isServerRunning.set(false); } @Test - @EnabledForJreRange(min = JAVA_16) public void testTimeout() throws Exception { TunnelingJdkSocket clientSocket = createClient(); InputStream inputStream = clientSocket.getInputStream(); @@ -155,12 +194,9 @@ public void testTimeout() throws Exception { clientSocket.close(); assertThrows(SocketException.class, () -> clientSocket.setSoTimeout(testTimeout)); assertThrows(SocketException.class, clientSocket::getSoTimeout); - - isServerRunning.set(false); } @Test - @EnabledForJreRange(min = JAVA_16) public void testBufferSizes() throws Exception { TunnelingJdkSocket clientSocket = createClient(); @@ -191,12 +227,9 @@ public void testBufferSizes() throws Exception { assertThrows(SocketException.class, clientSocket::getSendBufferSize); assertThrows(SocketException.class, clientSocket::getReceiveBufferSize); assertThrows(SocketException.class, clientSocket::getStreamBufferSize); - - isServerRunning.set(false); } @Test - @EnabledForJreRange(min = JAVA_16) public void testFileDescriptorLeak() throws Exception { long initialCount = getFileDescriptorCount(); @@ -209,15 +242,100 @@ public void testFileDescriptorLeak() throws Exception { } clientSocket.close(); - isServerRunning.set(false); + closeServer(); long finalCount = getFileDescriptorCount(); assertTrue(finalCount <= initialCount + 3); } + @Test + public void testClosedSelectorIsReportedAsSocketException() throws Exception { + try (TunnelingJdkSocket clientSocket = createClient()) { + InputStream inputStream = clientSocket.getInputStream(); + clientSocket.selector.close(); + + SocketException exception = assertThrows(SocketException.class, inputStream::read); + + assertInstanceOf(ClosedSelectorException.class, exception.getCause()); + } + } + + @Test + public void testSelectorClosedBetweenSelectAndSelectedKeysIsReportedAsSocketException() + throws Exception { + try (TunnelingJdkSocket clientSocket = createClient()) { + InputStream inputStream = clientSocket.getInputStream(); + clientSocket.selector.close(); + clientSocket.selector = new ClosedAfterSelectSelector(); + + SocketException exception = assertThrows(SocketException.class, inputStream::read); + + assertInstanceOf(ClosedSelectorException.class, exception.getCause()); + } + } + + @Test + public void testCancelledKeyIsReportedAsSocketException() throws Exception { + try (TunnelingJdkSocket clientSocket = createClient()) { + InputStream inputStream = clientSocket.getInputStream(); + clientSocket.selector.close(); + clientSocket.selector = new CancelledKeySelector(); + + SocketException exception = assertThrows(SocketException.class, inputStream::read); + + assertInstanceOf(CancelledKeyException.class, exception.getCause()); + } + } + + @Test + public void testAsynchronousCloseInterruptsBlockedReadWithIOException() throws Exception { + TunnelingJdkSocket clientSocket = createClient(); + Thread reader = null; + try { + InputStream inputStream = clientSocket.getInputStream(); + clientSocket.selector.close(); + BlockingCloseSelector selector = new BlockingCloseSelector(); + clientSocket.selector = selector; + AtomicReference readFailure = new AtomicReference<>(); + + reader = + new Thread( + () -> { + try { + inputStream.read(); + } catch (Throwable t) { + readFailure.set(t); + } + }, + "tunneling-jdk-socket-reader"); + reader.setDaemon(true); + reader.start(); + + assertTrue( + selector.awaitSelectStarted(5, TimeUnit.SECONDS), + "The reader did not block in Selector.select"); + clientSocket.close(); + reader.join(TimeUnit.SECONDS.toMillis(5)); + + assertFalse(reader.isAlive(), "The blocked read did not terminate after close"); + Throwable failure = readFailure.get(); + assertNotNull(failure, "The blocked read should fail when the socket is closed"); + assertTrue( + failure instanceof IOException, + () -> "Expected an IOException, but got " + failure.getClass().getName()); + assertFalse(failure instanceof ClosedSelectorException); + } finally { + clientSocket.close(); + if (reader != null && reader.isAlive()) { + reader.interrupt(); + reader.join(TimeUnit.SECONDS.toMillis(5)); + } + } + } + private long getFileDescriptorCount() { try { - Process process = Runtime.getRuntime().exec("lsof -p " + getPid()); + Process process = Runtime.getRuntime().exec("lsof -p " + ProcessHandle.current().pid()); int count = 0; try (java.io.BufferedReader reader = new java.io.BufferedReader(new java.io.InputStreamReader(process.getInputStream()))) { @@ -231,41 +349,8 @@ private long getFileDescriptorCount() { } } - private String getPid() { - return ManagementFactory.getRuntimeMXBean().getName().split("@")[0]; - } - - private static void startServer(UnixDomainSocketAddress socketAddress) { - Thread serverThread = - new Thread( - () -> { - try (ServerSocketChannel serverChannel = - ServerSocketChannel.open(StandardProtocolFamily.UNIX)) { - serverChannel.bind(socketAddress); - isServerRunning.set(true); - - synchronized (isServerRunning) { - isServerRunning.notifyAll(); - } - - while (isServerRunning.get()) { - SocketChannel clientChannel = serverChannel.accept(); - } - } catch (IOException e) { - throw new RuntimeException(e); - } - }); - serverThread.start(); - - synchronized (isServerRunning) { - while (!isServerRunning.get()) { - try { - isServerRunning.wait(); - } catch (InterruptedException e) { - throw new RuntimeException(e); - } - } - } + private void startServer(UnixDomainSocketAddress socketAddress) throws IOException { + server = new TestUnixSocketServer(socketAddress); } private Path getSocketPath() throws IOException { @@ -283,4 +368,279 @@ private TunnelingJdkSocket createClient() throws IOException { clientSocket.connect(new InetSocketAddress("localhost", 0)); return clientSocket; } + + private static final class TestUnixSocketServer implements AutoCloseable { + private final Path socketPath; + private final ServerSocketChannel serverChannel; + private final AtomicReference acceptedChannel = new AtomicReference<>(); + private final AtomicReference serverFailure = new AtomicReference<>(); + private final Thread serverThread; + + private TestUnixSocketServer(UnixDomainSocketAddress socketAddress) throws IOException { + socketPath = socketAddress.getPath(); + serverChannel = ServerSocketChannel.open(StandardProtocolFamily.UNIX); + boolean bound = false; + try { + serverChannel.bind(socketAddress); + bound = true; + } finally { + if (!bound) { + serverChannel.close(); + } + } + + serverThread = + new Thread( + () -> { + try { + acceptedChannel.set(serverChannel.accept()); + } catch (IOException e) { + if (serverChannel.isOpen()) { + serverFailure.set(e); + } + } + }, + "tunneling-jdk-socket-test-server"); + serverThread.setDaemon(true); + serverThread.start(); + } + + @Override + public void close() throws Exception { + serverChannel.close(); + serverThread.join(TimeUnit.SECONDS.toMillis(5)); + if (serverThread.isAlive()) { + serverThread.interrupt(); + throw new AssertionError("The Unix-domain test server did not terminate"); + } + + SocketChannel clientChannel = acceptedChannel.get(); + if (clientChannel != null) { + clientChannel.close(); + } + Files.deleteIfExists(socketPath); + + Throwable failure = serverFailure.get(); + if (failure != null) { + throw new AssertionError("The Unix-domain test server failed", failure); + } + } + } + + private abstract static class SelectorAdapter extends Selector { + private volatile boolean open = true; + + @Override + public final boolean isOpen() { + return open; + } + + @Override + public final SelectorProvider provider() { + return SelectorProvider.provider(); + } + + @Override + public final int selectNow() { + return doSelect(); + } + + @Override + public final int select(long timeout) { + return doSelect(); + } + + @Override + public final int select() { + return doSelect(); + } + + @Override + public final Selector wakeup() { + return this; + } + + @Override + public final void close() { + if (open) { + open = false; + onClose(); + } + } + + abstract int doSelect(); + + void onClose() {} + } + + /** + * Models another thread closing a selector while a socket read is blocked: + * + *
    + *
  1. Signals when {@code select()} is entered. + *
  2. Blocks until another thread closes the selector. + *
  3. Throws {@link ClosedSelectorException} after closure. + *
  4. Lets the test verify that the reader receives an {@link IOException} and terminates. + *
+ */ + private static final class BlockingCloseSelector extends SelectorAdapter { + private final CountDownLatch selectStarted = new CountDownLatch(1); + private final CountDownLatch closed = new CountDownLatch(1); + + @Override + public Set keys() { + return Collections.emptySet(); + } + + @Override + public Set selectedKeys() { + return Collections.emptySet(); + } + + @Override + int doSelect() { + selectStarted.countDown(); + try { + if (!closed.await(5, TimeUnit.SECONDS)) { + throw new AssertionError("Selector was not closed"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError("Interrupted while waiting for the selector to close", e); + } + throw new ClosedSelectorException(); + } + + @Override + void onClose() { + closed.countDown(); + } + + boolean awaitSelectStarted(long timeout, TimeUnit unit) throws InterruptedException { + return selectStarted.await(timeout, unit); + } + } + + /** + * Models a selector closing between selection and selected-key processing: + * + *
    + *
  1. Closes itself during {@code select()} and reports one ready channel. + *
  2. Throws {@link ClosedSelectorException} when the selected keys are requested. + *
  3. Lets the test verify that the exception is reported as a {@link SocketException}. + *
+ */ + private static final class ClosedAfterSelectSelector extends SelectorAdapter { + @Override + public Set keys() { + return Collections.emptySet(); + } + + @Override + public Set selectedKeys() { + if (!isOpen()) { + throw new ClosedSelectorException(); + } + return Collections.emptySet(); + } + + @Override + int doSelect() { + close(); + return 1; + } + } + + /** + * Models a key being cancelled before selected-key processing: + * + *
    + *
  1. Reports one ready channel from {@code select()}. + *
  2. Returns an invalid selected key. + *
  3. Throws {@link CancelledKeyException} when the read checks whether the key is readable. + *
  4. Lets the test verify that the exception is reported as a {@link SocketException}. + *
+ */ + private static final class CancelledKeySelector extends SelectorAdapter { + private final Set selectedKeys = + new HashSet<>(Collections.singleton(new CancelledSelectionKey(this))); + + @Override + public Set keys() { + return selectedKeys; + } + + @Override + public Set selectedKeys() { + return selectedKeys; + } + + @Override + int doSelect() { + return 1; + } + + private static final class CancelledSelectionKey extends SelectionKey { + private final Selector selector; + + private CancelledSelectionKey(Selector selector) { + this.selector = selector; + } + + @Override + public SelectableChannel channel() { + return null; + } + + @Override + public Selector selector() { + return selector; + } + + @Override + public boolean isValid() { + return false; + } + + @Override + public void cancel() {} + + @Override + public int interestOps() { + return OP_READ; + } + + @Override + public SelectionKey interestOps(int ops) { + return this; + } + + @Override + public int readyOps() { + throw new CancelledKeyException(); + } + } + } + + private static boolean udsSupported() { + Path socketPath = null; + try { + socketPath = Files.createTempFile("testSocketSupport", null); + Files.delete(socketPath); + try (ServerSocketChannel serverChannel = + ServerSocketChannel.open(StandardProtocolFamily.UNIX)) { + serverChannel.bind(UnixDomainSocketAddress.of(socketPath)); + } + return true; + } catch (IOException | UnsupportedOperationException e) { + return false; + } finally { + if (socketPath != null) { + try { + Files.deleteIfExists(socketPath); + } catch (IOException ignored) { + } + } + } + } }