Add option to enable jdk support for UDS

sarahchen6 · sarahchen6 · commit dbfc675c9c8b · 2025-03-26T17:58:14.000-04:00
diff --git a/src/main/java/com/timgroup/statsd/NonBlockingStatsDClient.java b/src/main/java/com/timgroup/statsd/NonBlockingStatsDClient.java
@@ -95,6 +95,7 @@ String tag() {
     public static final int SOCKET_BUFFER_BYTES = -1;
     public static final boolean DEFAULT_BLOCKING = false;
     public static final boolean DEFAULT_ENABLE_TELEMETRY = true;
+    public static final boolean DEFAULT_ENABLE_JDK_SOCKET = true;
 
     public static final boolean DEFAULT_ENABLE_AGGREGATION = true;
     public static final boolean DEFAULT_ENABLE_ORIGIN_DETECTION = true;
@@ -244,6 +245,9 @@ protected static String format(ThreadLocal<NumberFormat> formatter, Number value
      * @param connectionTimeout
      *     the timeout in milliseconds for connecting to the StatsD server. Applies to unix sockets only.
      *     It is also used to detect if a connection is still alive and re-establish a new one if needed.
+     * @param enableJdkSocket
+     *     Boolean to enable native JDK UDS support for the UnixStreamClientChannel.
+     *     Only compatible with Java 16 and up.
      * @throws StatsDClientException
      *     if the client could not be started
      */
@@ -253,7 +257,8 @@ protected static String format(ThreadLocal<NumberFormat> formatter, Number value
             final int maxPacketSizeBytes, String entityID, final int poolSize, final int processorWorkers,
             final int senderWorkers, boolean blocking, final boolean enableTelemetry, final int telemetryFlushInterval,
             final int aggregationFlushInterval, final int aggregationShards, final ThreadFactory customThreadFactory,
-            String containerID, final boolean originDetectionEnabled, final int connectionTimeout)
+            String containerID, final boolean originDetectionEnabled, final int connectionTimeout,
+            final boolean enableJdkSocket)
             throws StatsDClientException {
 
         if ((prefix != null) && (!prefix.isEmpty())) {
@@ -298,7 +303,7 @@ protected static String format(ThreadLocal<NumberFormat> formatter, Number value
         }
 
         try {
-            clientChannel = createByteChannel(addressLookup, timeout, connectionTimeout, bufferSize);
+            clientChannel = createByteChannel(addressLookup, timeout, connectionTimeout, bufferSize, enableJdkSocket);
 
             ThreadFactory threadFactory = customThreadFactory != null ? customThreadFactory : new StatsDThreadFactory();
 
@@ -317,7 +322,7 @@ protected static String format(ThreadLocal<NumberFormat> formatter, Number value
                 telemetryClientChannel = clientChannel;
                 telemetryStatsDProcessor = statsDProcessor;
             } else {
-                telemetryClientChannel = createByteChannel(telemetryAddressLookup, timeout, connectionTimeout, bufferSize);
+                telemetryClientChannel = createByteChannel(telemetryAddressLookup, timeout, connectionTimeout, bufferSize, enableJdkSocket);
 
                 // similar settings, but a single worker and non-blocking.
                 telemetryStatsDProcessor = createProcessor(queueSize, handler, getPacketSize(telemetryClientChannel),
@@ -378,7 +383,8 @@ public NonBlockingStatsDClient(final NonBlockingStatsDClientBuilder builder) thr
             builder.blocking, builder.enableTelemetry, builder.telemetryFlushInterval,
             (builder.enableAggregation ? builder.aggregationFlushInterval : 0),
             builder.aggregationShards, builder.threadFactory, builder.containerID,
-            builder.originDetectionEnabled, builder.connectionTimeout);
+            builder.originDetectionEnabled, builder.connectionTimeout,
+            builder.enableJdkSocket);
     }
 
     protected StatsDProcessor createProcessor(final int queueSize, final StatsDClientErrorHandler handler,
@@ -480,7 +486,7 @@ StringBuilder tagString(final String[] tags, StringBuilder builder) {
     }
 
     ClientChannel createByteChannel(
-            Callable<SocketAddress> addressLookup, int timeout, int connectionTimeout, int bufferSize)
+            Callable<SocketAddress> addressLookup, int timeout, int connectionTimeout, int bufferSize, boolean enableJdkSocket)
             throws Exception {
         final SocketAddress address = addressLookup.call();
         if (address instanceof NamedPipeSocketAddress) {
@@ -493,7 +499,7 @@ ClientChannel createByteChannel(
             // Allow us to support `unix://` for both kind of sockets like in go.
             switch (unixAddr.getTransportType()) {
                 case UDS_STREAM:
-                    return new UnixStreamClientChannel(unixAddr.getAddress(), timeout, connectionTimeout, bufferSize);
+                    return new UnixStreamClientChannel(unixAddr.getAddress(), timeout, connectionTimeout, bufferSize, enableJdkSocket);
                 case UDS_DATAGRAM:
                 case UDS:
                     return new UnixDatagramClientChannel(unixAddr.getAddress(), timeout, bufferSize);
diff --git a/src/main/java/com/timgroup/statsd/NonBlockingStatsDClientBuilder.java b/src/main/java/com/timgroup/statsd/NonBlockingStatsDClientBuilder.java
@@ -33,6 +33,7 @@ public class NonBlockingStatsDClientBuilder implements Cloneable {
     public boolean blocking = NonBlockingStatsDClient.DEFAULT_BLOCKING;
     public boolean enableTelemetry = NonBlockingStatsDClient.DEFAULT_ENABLE_TELEMETRY;
     public boolean enableAggregation = NonBlockingStatsDClient.DEFAULT_ENABLE_AGGREGATION;
+    public boolean enableJdkSocket = NonBlockingStatsDClient.DEFAULT_ENABLE_JDK_SOCKET;
     public int telemetryFlushInterval = Telemetry.DEFAULT_FLUSH_INTERVAL;
     public int aggregationFlushInterval = StatsDAggregator.DEFAULT_FLUSH_INTERVAL;
     public int aggregationShards = StatsDAggregator.DEFAULT_SHARDS;
@@ -205,6 +206,11 @@ public NonBlockingStatsDClientBuilder originDetectionEnabled(boolean val) {
         return this;
     }
 
+    public NonBlockingStatsDClientBuilder enableJdkSocket(boolean val) {
+        enableJdkSocket = val;
+        return this;
+    }
+
     /**
      * NonBlockingStatsDClient factory method.
      * @return the built NonBlockingStatsDClient.
@@ -380,7 +386,7 @@ protected static Callable<SocketAddress> staticUnixResolution(
             @Override public SocketAddress call() {
                 SocketAddress socketAddress;
                 // Use native UDS support for compatible Java versions and jnr-unixsocket support otherwise.
-                if (VersionUtils.isJavaVersionAtLeast(16)) {
+                if (VersionUtils.isJavaVersionAtLeast(16) && NonBlockingStatsDClient.DEFAULT_ENABLE_JDK_SOCKET) {
                     try {
                         // Use reflection to avoid compiling Java 16+ classes in incompatible versions
                         Class<?> unixDomainSocketAddressClass = Class.forName("java.net.UnixDomainSocketAddress");
diff --git a/src/main/java/com/timgroup/statsd/UnixStreamClientChannel.java b/src/main/java/com/timgroup/statsd/UnixStreamClientChannel.java
@@ -20,6 +20,7 @@ public class UnixStreamClientChannel implements ClientChannel {
     private final int timeout;
     private final int connectionTimeout;
     private final int bufferSize;
+    private final boolean enableJdkSocket;
 
     private SocketChannel delegate;
     private final ByteBuffer delimiterBuffer = ByteBuffer.allocateDirect(Integer.SIZE / Byte.SIZE).order(ByteOrder.LITTLE_ENDIAN);
@@ -29,14 +30,15 @@ public class UnixStreamClientChannel implements ClientChannel {
      *
      * @param address Location of named pipe
      */
-    UnixStreamClientChannel(SocketAddress address, int timeout, int connectionTimeout, int bufferSize) throws IOException {
+    UnixStreamClientChannel(SocketAddress address, int timeout, int connectionTimeout, int bufferSize, boolean enableJdkSocket) throws IOException {
         this.delegate = null;
         this.address = address;
         System.out.println("========== Constructor address: " + address);
         System.out.println("========== Constructor address type: " + address.getClass().getName());
         this.timeout = timeout;
         this.connectionTimeout = connectionTimeout;
         this.bufferSize = bufferSize;
+        this.enableJdkSocket = enableJdkSocket;
     }
 
     @Override
@@ -130,7 +132,7 @@ private void connect() throws IOException {
 
         long deadline = System.nanoTime() + connectionTimeout * 1_000_000L;
         // Use native UDS support for compatible Java versions and jnr-unixsocket support otherwise.
-        if (VersionUtils.isJavaVersionAtLeast(16)) {
+        if (VersionUtils.isJavaVersionAtLeast(16) && enableJdkSocket) {
             try {
                 // Use reflection to avoid compiling Java 16+ classes in incompatible versions
                 Class<?> protocolFamilyClass = Class.forName("java.net.StandardProtocolFamily");
@@ -177,7 +179,7 @@ private void connect() throws IOException {
                 throw new IOException("Failed to create UnixStreamClientChannel for native UDS implementation", e);
             }
         }
-        // Default to jnr-unixsocket if Java version is less than 16
+        // Default to jnr-unixsocket if Java version is less than 16 or native UDS support is disabled
         UnixSocketChannel channel = UnixSocketChannel.create();
         
         if (connectionTimeout > 0) {
diff --git a/src/test/java/com/timgroup/statsd/NonBlockingStatsDClientTest.java b/src/test/java/com/timgroup/statsd/NonBlockingStatsDClientTest.java
@@ -1831,7 +1831,7 @@ public NonBlockingStatsDClient build() {
                     this.originDetectionEnabled(false);
                     return new NonBlockingStatsDClient(resolve()) {
                         @Override
-                        ClientChannel createByteChannel(Callable<SocketAddress> addressLookup, int timeout, int connectionTimeout, int bufferSize) throws Exception {
+                        ClientChannel createByteChannel(Callable<SocketAddress> addressLookup, int timeout, int connectionTimeout, int bufferSize, boolean enableJdkSocket) throws Exception {
                             return new DatagramClientChannel(addressLookup.call()) {
                                 @Override
                                 public int write(ByteBuffer data) throws IOException {
@@ -1870,7 +1870,7 @@ public NonBlockingStatsDClient build() {
                     this.bufferPoolSize(1);
                     return new NonBlockingStatsDClient(resolve()) {
                         @Override
-                        ClientChannel createByteChannel(Callable<SocketAddress> addressLookup, int timeout, int connectionTimeout, int bufferSize) throws Exception {
+                        ClientChannel createByteChannel(Callable<SocketAddress> addressLookup, int timeout, int connectionTimeout, int bufferSize, boolean enableJdkSocket) throws Exception {
                             return new DatagramClientChannel(addressLookup.call()) {
                                 @Override
                                 public int write(ByteBuffer data) throws IOException {
