From b8dcf102cd8d58fce7ac898cf8652b8012124597 Mon Sep 17 00:00:00 2001 From: Brice Dutheil Date: Wed, 26 Aug 2026 18:29:39 +0200 Subject: [PATCH 1/3] fix: Harden TunnelingUnixSocket handling of bad state --- utils/socket-utils/build.gradle.kts | 1 + .../common/socket/TunnelingJdkSocket.java | 119 +++-- .../common/socket/TunnelingJdkSocketTest.java | 452 +++++++++++++++--- 3 files changed, 466 insertions(+), 106 deletions(-) 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/java17/datadog/common/socket/TunnelingJdkSocket.java b/utils/socket-utils/src/main/java17/datadog/common/socket/TunnelingJdkSocket.java index fb25228e0f6..b017cf32322 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; @@ -10,6 +11,8 @@ import java.net.SocketException; 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; @@ -28,13 +31,13 @@ final class TunnelingJdkSocket extends Socket { private final SocketAddress unixSocketAddress; private InetSocketAddress inetSocketAddress; - private SocketChannel unixSocketChannel; - private Selector selector; + private volatile SocketChannel unixSocketChannel; + @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 @@ -90,7 +93,7 @@ public synchronized int getSoTimeout() throws SocketException { } @Override - public void connect(final SocketAddress endpoint) throws IOException { + public synchronized void connect(final SocketAddress endpoint) throws IOException { if (endpoint == null) { throw new IllegalArgumentException("Endpoint cannot be null"); } @@ -108,7 +111,8 @@ public void connect(final SocketAddress endpoint) throws IOException { // See: // https://github.com/jnr/jnr-unixsocket/blob/master/src/main/java/jnr/unixsocket/UnixSocket.java#L89-L97 @Override - public void connect(final SocketAddress endpoint, final int timeout) throws IOException { + public synchronized void connect(final SocketAddress endpoint, final int timeout) + throws IOException { if (endpoint == null) { throw new IllegalArgumentException("Endpoint cannot be null"); } @@ -199,7 +203,7 @@ public int getStreamBufferSize() throws SocketException { } @Override - public InputStream getInputStream() throws IOException { + public synchronized InputStream getInputStream() throws IOException { if (isClosed()) { throw new SocketException("Socket is closed"); } @@ -210,10 +214,21 @@ public InputStream getInputStream() throws IOException { throw new SocketException("Socket input is shutdown"); } - if (selector == null) { - selector = Selector.open(); - unixSocketChannel.configureBlocking(false); - unixSocketChannel.register(selector, SelectionKey.OP_READ); + 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; + } } return new InputStream() { @@ -232,30 +247,45 @@ public int read(byte[] b, int off, int len) throws IOException { } buffer.clear(); - int readyChannels = selector.select(timeout); - if (readyChannels == 0) { - return 0; + Selector currentSelector = selector; + SocketChannel currentChannel = unixSocketChannel; + if (currentSelector == null || currentChannel == null) { + throw new SocketException("Socket is closed"); } - 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; + try { + int readyChannels = currentSelector.select(timeout); + if (readyChannels == 0) { + if (isClosed() || !currentSelector.isOpen()) { + throw new SocketException("Socket is closed"); + } + return 0; + } + + Set selectedKeys = currentSelector.selectedKeys(); + synchronized (selectedKeys) { + Iterator keyIterator = selectedKeys.iterator(); + while (keyIterator.hasNext()) { + SelectionKey key = keyIterator.next(); + keyIterator.remove(); + if (key.isReadable()) { + int r = currentChannel.read(buffer); + if (r == -1) { + return -1; + } + buffer.flip(); + len = Math.min(r, len); + buffer.get(b, off, len); + return len; } - 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 @@ -303,7 +333,7 @@ public void close() throws IOException { } @Override - public void shutdownInput() throws IOException { + public synchronized void shutdownInput() throws IOException { if (isClosed()) { throw new SocketException("Socket is closed"); } @@ -318,7 +348,7 @@ public void shutdownInput() throws IOException { } @Override - public void shutdownOutput() throws IOException { + public synchronized void shutdownOutput() throws IOException { if (isClosed()) { throw new SocketException("Socket is closed"); } @@ -341,36 +371,25 @@ public InetAddress getInetAddress() { } @Override - public void close() throws IOException { + public synchronized void close() { if (isClosed()) { return; } + closed = true; + shutIn = true; + shutOut = true; // Ignore possible exceptions so that we continue closing the socket - try { - if (!isInputShutdown()) { - shutdownInput(); - } - } catch (IOException e) { - } - try { - if (!isOutputShutdown()) { - shutdownOutput(); - } - } catch (IOException e) { - } try { if (selector != null) { selector.close(); - selector = null; } - } catch (IOException e) { + } catch (IOException ignored) { } try { if (unixSocketChannel != null) { unixSocketChannel.close(); } - } catch (IOException e) { + } 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..fb930e3a928 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,26 +13,52 @@ 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); @@ -62,12 +90,9 @@ public void testSocketConnectAndClose() throws Exception { assertThrows(SocketException.class, clientSocket::getInputStream); assertThrows(SocketException.class, clientSocket::getOutputStream); clientSocket.close(); - - isServerRunning.set(false); } @Test - @EnabledForJreRange(min = JAVA_16) public void testInputStreamClose() throws Exception { TunnelingJdkSocket clientSocket = createClient(); InputStream inputStream = clientSocket.getInputStream(); @@ -86,12 +111,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 +132,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 +174,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 +207,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 +222,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 +329,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 +348,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) { + } + } + } + } } From 252b8b5e8ad2539d0c66c7976425a5f95fa1fb35 Mon Sep 17 00:00:00 2001 From: Brice Dutheil Date: Thu, 27 Aug 2026 17:25:47 +0200 Subject: [PATCH 2/3] fix: make JDK UDS lifecycle race-safe Create the `SocketChannel` with `TunnelingJdkSocket` instead of waiting until `connect` so they share the same lifetime. This gives each socket one stable channel and ensures a failed connection closes all associated resources. Okio may close a socket from its _timeout watchdog_ while another thread is _initializing the read selector_ or _performing a half-close_. Coordinate these compound lifecycle operations on the socket monitor: - `getInputStream` publishes the selector before `close` can snapshot it, also it runs once per physical OkHttp connection - `shutdownInput` and `shutdownOutput` complete their channel operation and state update without interleaving with `close`. Those are lifecycle operations. - `close` publishes the terminal state and snapshots the selector atomically, then releases the monitor before closing resources and waking blocked reads. This is a lifecycle operation. - The normal bulk reads and writes only perform volatile state reads. - Previous synchronization on the selection key is unchanged. So, this introduces the socket monitor only for these lifecycle transitions. This keeps normal reads and writes outside the socket monitor. Only selector publication and socket lifecycle transitions require that coordination. --- .../socket/UnixDomainSocketFactory.java | 2 +- .../common/socket/TunnelingJdkSocket.java | 275 +++++++++--------- .../common/socket/TunnelingJdkSocketTest.java | 20 ++ 3 files changed, 162 insertions(+), 135 deletions(-) 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 b017cf32322..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 @@ -9,6 +9,7 @@ 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; @@ -28,10 +29,10 @@ * 16. */ final class TunnelingJdkSocket extends Socket { - private final SocketAddress unixSocketAddress; - private InetSocketAddress inetSocketAddress; + private final UnixDomainSocketAddress unixSocketAddress; + private final SocketChannel unixSocketChannel; - private volatile SocketChannel unixSocketChannel; + private volatile InetSocketAddress inetSocketAddress; @VisibleForTesting volatile Selector selector; private volatile int timeout; @@ -44,18 +45,14 @@ final class TunnelingJdkSocket extends Socket { 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 @@ -74,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"); } @@ -85,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"); } @@ -93,26 +90,15 @@ public synchronized int getSoTimeout() throws SocketException { } @Override - public synchronized 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); + public void connect(final SocketAddress endpoint) throws IOException { + connect(endpoint, 0); } // `timeout` is intentionally ignored here, like in the jnr-unixsocket implementation. // See: // https://github.com/jnr/jnr-unixsocket/blob/master/src/main/java/jnr/unixsocket/UnixSocket.java#L89-L97 @Override - public synchronized void connect(final SocketAddress endpoint, final int timeout) - throws IOException { + public void connect(final SocketAddress endpoint, final int timeout) throws IOException { if (endpoint == null) { throw new IllegalArgumentException("Endpoint cannot be null"); } @@ -125,8 +111,14 @@ public synchronized void connect(final SocketAddress endpoint, final int timeout 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 @@ -203,96 +195,101 @@ public int getStreamBufferSize() throws SocketException { } @Override - public synchronized 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"); - } + public InputStream getInputStream() throws IOException { + // 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"); + } - 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) { + Selector currentSelector = selector; + if (currentSelector == null) { + currentSelector = Selector.open(); try { - currentSelector.close(); - } catch (IOException closeException) { - e.addSuppressed(closeException); + 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; } - throw e; } - } - return new InputStream() { - private final ByteBuffer buffer = ByteBuffer.allocate(getStreamBufferSize()); + 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); - } - - @Override - public int read(byte[] b, int off, int len) throws IOException { - if (isInputShutdown()) { - return -1; + @Override + public int read() throws IOException { + byte[] nextByte = new byte[1]; + return (read(nextByte, 0, 1) == -1) ? -1 : (nextByte[0] & 0xFF); } - buffer.clear(); - Selector currentSelector = selector; - SocketChannel currentChannel = unixSocketChannel; - if (currentSelector == null || currentChannel == null) { - throw new SocketException("Socket is closed"); - } + @Override + public int read(byte[] b, int off, int len) throws IOException { + if (isInputShutdown()) { + return -1; + } + buffer.clear(); - try { - int readyChannels = currentSelector.select(timeout); - if (readyChannels == 0) { - if (isClosed() || !currentSelector.isOpen()) { - throw new SocketException("Socket is closed"); - } - return 0; + Selector currentSelector = selector; + if (currentSelector == null) { + throw new SocketException("Socket is closed"); } - Set selectedKeys = currentSelector.selectedKeys(); - synchronized (selectedKeys) { - Iterator keyIterator = selectedKeys.iterator(); - while (keyIterator.hasNext()) { - SelectionKey key = keyIterator.next(); - keyIterator.remove(); - if (key.isReadable()) { - int r = currentChannel.read(buffer); - if (r == -1) { - return -1; + try { + int readyChannels = currentSelector.select(timeout); + if (readyChannels == 0) { + if (isClosed() || !currentSelector.isOpen()) { + throw new SocketException("Socket is closed"); + } + 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; } - 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; - } catch (ClosedSelectorException | CancelledKeyException e) { - SocketException socketException = new SocketException("Socket is closed"); - socketException.initCause(e); - throw socketException; } - } - @Override - public void close() throws IOException { - TunnelingJdkSocket.this.close(); - } - }; + @Override + public void close() throws IOException { + TunnelingJdkSocket.this.close(); + } + }; + } } @Override @@ -333,33 +330,39 @@ public void close() throws IOException { } @Override - public synchronized 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"); + public void shutdownInput() throws IOException { + // 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 synchronized 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"); + public void shutdownOutput() throws IOException { + // 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 @@ -371,24 +374,28 @@ public InetAddress getInetAddress() { } @Override - public synchronized void close() { - if (isClosed()) { - return; + 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; + } + shutIn = true; + shutOut = true; + closed = true; + currentSelector = selector; } - closed = true; - shutIn = true; - shutOut = true; // Ignore possible exceptions so that we continue closing the socket try { - if (selector != null) { - selector.close(); + if (currentSelector != null) { + currentSelector.close(); } } catch (IOException ignored) { } try { - if (unixSocketChannel != null) { - unixSocketChannel.close(); - } + unixSocketChannel.close(); } catch (IOException ignored) { } } 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 fb930e3a928..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 @@ -65,6 +65,8 @@ public void testSocketConnectAndClose() throws Exception { startServer(socketAddress); TunnelingJdkSocket clientSocket = new TunnelingJdkSocket(socketPath); + assertNotNull(clientSocket.getChannel()); + assertTrue(clientSocket.getChannel().isOpen()); assertFalse(clientSocket.isConnected()); assertFalse(clientSocket.isClosed()); @@ -83,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()); @@ -92,6 +95,23 @@ public void testSocketConnectAndClose() throws Exception { 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))); + + 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 public void testInputStreamClose() throws Exception { TunnelingJdkSocket clientSocket = createClient(); From 39570c8c95822ecb02060adfae1c6a4581f85bfa Mon Sep 17 00:00:00 2001 From: Brice Dutheil Date: Thu, 27 Aug 2026 17:51:56 +0200 Subject: [PATCH 3/3] chore: Adds a behavioral test on the DDAgentWriter when UDS is closed --- dd-trace-core/build.gradle | 2 + .../socket/UnixDomainServerSocketFactory.java | 109 ++++++++++++++++++ .../writer/DDAgentWriterCombinedTest.java | 108 +++++++++++++++-- 3 files changed, 210 insertions(+), 9 deletions(-) create mode 100644 dd-trace-core/src/test/java/datadog/common/socket/UnixDomainServerSocketFactory.java 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); } }