Merge pull request #228 from iksaif/corentin.chary/unix-stream

Author: iksaif

CHANGELOG.md

@@ -1,5 +1,9 @@
 # CHANGELOG
 
+## 4.3.0 / 2024.XX.XX
+
+* [FEATURE] Add support for `SOCK_STREAM` Unix sockets. See [#228][]
+
 ## 4.2.1 / 2023.03.10
 
 * [FEATURE] Add support for `DD_DOGSTATSD_URL`. See [#217][]
@@ -232,6 +236,7 @@ Fork from [indeedeng/java-dogstatsd-client] (https://github.com/indeedeng/java-d
 [#203]: https://github.com/DataDog/java-dogstatsd-client/issues/203
 [#211]: https://github.com/DataDog/java-dogstatsd-client/issues/211
 [#217]: https://github.com/DataDog/java-dogstatsd-client/issues/217
+[#228]: https://github.com/DataDog/java-dogstatsd-client/pull/228
 
 [@PatrickAuld]: https://github.com/PatrickAuld
 [@blevz]: https://github.com/blevz

README.md

@@ -23,13 +23,23 @@ The client jar is distributed via Maven central, and can be downloaded [from Mav
 
 ### Unix Domain Socket support
 
-As an alternative to UDP, Agent v6 can receive metrics via a UNIX Socket (on Linux only). This library supports transmission via this protocol. To use it, pass the socket path as a hostname, and `0` as port.
+As an alternative to UDP, Agent v6 can receive metrics via a UNIX Socket (on Linux only). This library supports transmission via this protocol. To use it
+use the `address()` method of the builder and pass the path to the socket with the `unix://` prefix:
+
+```java
+StatsDClient client = new NonBlockingStatsDClientBuilder()
+    .address("unix:///var/run/datadog/dsd.socket")
+    .build();
+```
 
 By default, all exceptions are ignored, mimicking UDP behaviour. When using Unix Sockets, transmission errors trigger exceptions you can choose to handle by passing a `StatsDClientErrorHandler`:
 
 - Connection error because of an invalid/missing socket triggers a `java.io.IOException: No such file or directory`.
 - If DogStatsD's reception buffer were to fill up and the non blocking client is used, the send times out after 100ms and throw either a `java.io.IOException: No buffer space available` or a `java.io.IOException: Resource temporarily unavailable`.
 
+The default UDS transport is using `SOCK_DATAGRAM` sockets. We also have experimental support for `SOCK_STREAM` sockets which can
+be enabled by using the `unixstream://` instead of `unix://`. This is not recommended for production use at this time.
+
 ## Configuration
 
 Once your DogStatsD client is installed, instantiate it in your code:

src/main/java/com/timgroup/statsd/NonBlockingStatsDClient.java

@@ -18,7 +18,6 @@
 import java.util.concurrent.Callable;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.ThreadLocalRandom;
-import java.util.concurrent.TimeUnit;
 
 
 /**
@@ -99,6 +98,7 @@ String tag() {
 
     public static final boolean DEFAULT_ENABLE_AGGREGATION = true;
     public static final boolean DEFAULT_ENABLE_ORIGIN_DETECTION = true;
+    public static final int SOCKET_CONNECT_TIMEOUT_MS = 1000;
 
     public static final String CLIENT_TAG = "client:java";
     public static final String CLIENT_VERSION_TAG = "client_version:";
@@ -241,6 +241,9 @@ protected static String format(ThreadLocal<NumberFormat> formatter, Number value
      *     The client tries to read the container ID by parsing the file /proc/self/cgroup.
      *     This is not supported on Windows.
      *     The client prioritizes the value passed via or entityID or DD_ENTITY_ID (if set) over the container ID.
+     * @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.
      * @throws StatsDClientException
      *     if the client could not be started
      */
@@ -250,7 +253,7 @@ private NonBlockingStatsDClient(final String prefix, final int queueSize, final
             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)
+            String containerID, final boolean originDetectionEnabled, final int connectionTimeout)
             throws StatsDClientException {
 
         if ((prefix != null) && (!prefix.isEmpty())) {
@@ -297,7 +300,7 @@ private NonBlockingStatsDClient(final String prefix, final int queueSize, final
         }
 
         try {
-            clientChannel = createByteChannel(addressLookup, timeout, bufferSize);
+            clientChannel = createByteChannel(addressLookup, timeout, connectionTimeout, bufferSize);
 
             ThreadFactory threadFactory = customThreadFactory != null ? customThreadFactory : new StatsDThreadFactory();
 
@@ -316,7 +319,7 @@ private NonBlockingStatsDClient(final String prefix, final int queueSize, final
                 telemetryClientChannel = clientChannel;
                 telemetryStatsDProcessor = statsDProcessor;
             } else {
-                telemetryClientChannel = createByteChannel(telemetryAddressLookup, timeout, bufferSize);
+                telemetryClientChannel = createByteChannel(telemetryAddressLookup, timeout, connectionTimeout, bufferSize);
 
                 // similar settings, but a single worker and non-blocking.
                 telemetryStatsDProcessor = createProcessor(queueSize, handler, maxPacketSizeBytes,
@@ -377,7 +380,7 @@ 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.originDetectionEnabled, builder.connectionTimeout);
     }
 
     protected StatsDProcessor createProcessor(final int queueSize, final StatsDClientErrorHandler handler,
@@ -478,11 +481,29 @@ StringBuilder tagString(final String[] tags, StringBuilder builder) {
         return tagString(tags, constantTagsRendered, builder);
     }
 
-    ClientChannel createByteChannel(Callable<SocketAddress> addressLookup, int timeout, int bufferSize) throws Exception {
+    ClientChannel createByteChannel(
+            Callable<SocketAddress> addressLookup, int timeout, int connectionTimeout, int bufferSize)
+            throws Exception {
         final SocketAddress address = addressLookup.call();
         if (address instanceof NamedPipeSocketAddress) {
             return new NamedPipeClientChannel((NamedPipeSocketAddress) address);
         }
+        if (address instanceof UnixSocketAddressWithTransport) {
+            UnixSocketAddressWithTransport unixAddr = ((UnixSocketAddressWithTransport) address);
+
+            // TODO: Maybe introduce a `UnixClientChannel` that can handle both stream and datagram sockets? This would
+            // 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);
+                case UDS_DATAGRAM:
+                case UDS:
+                    return new UnixDatagramClientChannel(unixAddr.getAddress(), timeout, bufferSize);
+                default:
+                    throw new IllegalArgumentException("Unsupported transport type: " + unixAddr.getTransportType());
+            }
+        }
+        // We keep this for backward compatibility
         try {
             if (Class.forName("jnr.unixsocket.UnixSocketAddress").isInstance(address)) {
                 return new UnixDatagramClientChannel(address, timeout, bufferSize);

src/main/java/com/timgroup/statsd/NonBlockingStatsDClientBuilder.java

@@ -1,5 +1,6 @@
 package com.timgroup.statsd;
 
+import jnr.constants.platform.Sock;
 import jnr.unixsocket.UnixSocketAddress;
 
 import java.net.InetAddress;
@@ -34,6 +35,7 @@ public class NonBlockingStatsDClientBuilder implements Cloneable {
     public int aggregationFlushInterval = StatsDAggregator.DEFAULT_FLUSH_INTERVAL;
     public int aggregationShards = StatsDAggregator.DEFAULT_SHARDS;
     public boolean originDetectionEnabled = NonBlockingStatsDClient.DEFAULT_ENABLE_ORIGIN_DETECTION;
+    public int connectionTimeout = NonBlockingStatsDClient.SOCKET_CONNECT_TIMEOUT_MS;
 
     public Callable<SocketAddress> addressLookup;
     public Callable<SocketAddress> telemetryAddressLookup;
@@ -71,6 +73,11 @@ public NonBlockingStatsDClientBuilder timeout(int val) {
         return this;
     }
 
+    public NonBlockingStatsDClientBuilder connectionTimeout(int val) {
+        connectionTimeout = val;
+        return this;
+    }
+
     public NonBlockingStatsDClientBuilder bufferPoolSize(int val) {
         bufferPoolSize = val;
         return this;
@@ -126,6 +133,16 @@ public NonBlockingStatsDClientBuilder namedPipe(String val) {
         return this;
     }
 
+    public NonBlockingStatsDClientBuilder address(String address) {
+        addressLookup = getAddressLookupFromUrl(address);
+        return this;
+    }
+
+    public NonBlockingStatsDClientBuilder telemetryAddress(String address) {
+        telemetryAddressLookup = getAddressLookupFromUrl(address);
+        return this;
+    }
+
     public NonBlockingStatsDClientBuilder prefix(String val) {
         prefix = val;
         return this;
@@ -283,9 +300,12 @@ private Callable<SocketAddress> getAddressLookupFromUrl(String url) {
             return staticAddress(uriHost, uriPort);
         }
 
-        if (parsed.getScheme().equals("unix")) {
+        if (parsed.getScheme().startsWith("unix")) {
             String uriPath = parsed.getPath();
-            return staticAddress(uriPath, 0);
+            return staticUnixResolution(
+                    uriPath,
+                    UnixSocketAddressWithTransport.TransportType.fromScheme(parsed.getScheme())
+            );
         }
 
         return null;
@@ -304,7 +324,10 @@ public static Callable<SocketAddress> volatileAddressResolution(final String hos
         if (port == 0) {
             return new Callable<SocketAddress>() {
                 @Override public SocketAddress call() throws UnknownHostException {
-                    return new UnixSocketAddress(hostname);
+                    return new UnixSocketAddressWithTransport(
+                            new UnixSocketAddress(hostname),
+                            UnixSocketAddressWithTransport.TransportType.UDS
+                    );
                 }
             };
         } else {
@@ -343,6 +366,17 @@ protected static Callable<SocketAddress> staticNamedPipeResolution(String namedP
         };
     }
 
+    protected static Callable<SocketAddress> staticUnixResolution(
+            final String path,
+            final UnixSocketAddressWithTransport.TransportType transportType) {
+        return new Callable<SocketAddress>() {
+            @Override public SocketAddress call() {
+                final UnixSocketAddress socketAddress = new UnixSocketAddress(path);
+                return new UnixSocketAddressWithTransport(socketAddress, transportType);
+            }
+        };
+    }
+
     private static Callable<SocketAddress> staticAddress(final String hostname, final int port) {
         try {
             return staticAddressResolution(hostname, port);

src/main/java/com/timgroup/statsd/UnixSocketAddressWithTransport.java

@@ -0,0 +1,70 @@
+package com.timgroup.statsd;
+
+import java.net.SocketAddress;
+import java.util.Objects;
+
+public class UnixSocketAddressWithTransport extends SocketAddress {
+
+    private final SocketAddress address;
+    private final TransportType transportType;
+
+    public enum TransportType {
+        UDS_STREAM("uds-stream"),
+        UDS_DATAGRAM("uds-datagram"),
+        UDS("uds");
+
+        private final String transportType;
+
+        TransportType(String transportType) {
+            this.transportType = transportType;
+        }
+
+        String getTransportType() {
+            return transportType;
+        }
+
+        static TransportType fromScheme(String scheme) {
+            switch (scheme) {
+                case "unixstream":
+                    return UDS_STREAM;
+                case "unixgram":
+                    return UDS_DATAGRAM;
+                case "unix":
+                    return UDS;
+                default:
+                    break;
+            }
+            throw new IllegalArgumentException("Unknown scheme: " + scheme);
+        }
+    }
+
+    public UnixSocketAddressWithTransport(final SocketAddress address, final TransportType transportType) {
+        this.address = address;
+        this.transportType = transportType;
+    }
+
+    @Override
+    public boolean equals(Object other) {
+        if (this == other) {
+            return true;
+        }
+        if (other == null || getClass() != other.getClass()) {
+            return false;
+        }
+        UnixSocketAddressWithTransport that = (UnixSocketAddressWithTransport) other;
+        return Objects.equals(address, that.address) && transportType == that.transportType;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(address, transportType);
+    }
+
+    SocketAddress getAddress() {
+        return address;
+    }
+
+    TransportType getTransportType() {
+        return transportType;
+    }
+}