Add builder method

DataDog · Jan 17, 2024 · d50bb85 · d50bb85
1 parent 95843b6
commit d50bb85
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 17 deletions.
diff --git a/src/main/java/com/timgroup/statsd/NonBlockingStatsDClientBuilder.java b/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;
@@ -132,6 +133,25 @@ public NonBlockingStatsDClientBuilder namedPipe(String val) {
         return this;
     }
 
+    private Callable<SocketAddress> socketLookup(final String path, final UnixSocketAddressWithTransport.TransportType transport) {
+        return new Callable<SocketAddress>() {
+            @Override
+            public SocketAddress call() throws Exception {
+                return new UnixSocketAddressWithTransport(new UnixSocketAddress(path), transport);
+            }
+        };
+    }
+
+    public NonBlockingStatsDClientBuilder socket(final String path, final UnixSocketAddressWithTransport.TransportType transport) {
+        addressLookup = socketLookup(path, transport);
+        return this;
+    }
+
+    public NonBlockingStatsDClientBuilder telemetrySocket(final String path, final UnixSocketAddressWithTransport.TransportType transport) {
+        telemetryAddressLookup = socketLookup(path, transport);
+        return this;
+    }
+
     public NonBlockingStatsDClientBuilder prefix(String val) {
         prefix = val;
         return this;

diff --git a/src/main/java/com/timgroup/statsd/UnixStreamClientChannel.java b/src/main/java/com/timgroup/statsd/UnixStreamClientChannel.java
@@ -134,17 +134,13 @@ private void connect() throws IOException {
             // We'd have better timeout support if we used Java 16's native Unix domain socket support (JEP 380)
             delegate.setOption(UnixSocketOptions.SO_SNDTIMEO, connectionTimeout);
         }
-        delegate.connect(address);
-        while (!delegate.finishConnect()) {
-            // wait for connection to be established
-            try {
-                Thread.sleep(10);
-            } catch (InterruptedException e) {
-                throw new IOException("Interrupted while waiting for connection", e);
-            }
+        if (!delegate.connect(address)) {
             if (connectionTimeout > 0 && System.nanoTime() > deadline) {
                 throw new IOException("Connection timed out");
             }
+            if (!delegate.finishConnect()) {
+                throw new IOException("Connection failed");
+            }
         }
 
         if (timeout > 0) {

diff --git a/src/test/java/com/timgroup/statsd/UnixStreamSocketTest.java b/src/test/java/com/timgroup/statsd/UnixStreamSocketTest.java
@@ -51,15 +51,8 @@ public void start() throws IOException {
 
         server = new UnixStreamSocketDummyStatsDServer(socketFile.toString());
 
-        Callable<SocketAddress> addressLookup = new Callable<SocketAddress>() {
-            @Override
-            public SocketAddress call() throws Exception {
-                return new UnixSocketAddressWithTransport(new UnixSocketAddress(socketFile.getPath()), UnixSocketAddressWithTransport.TransportType.UDS_STREAM);
-            }
-        };
-
         client = new NonBlockingStatsDClientBuilder().prefix("my.prefix")
-            .addressLookup(addressLookup)
+            .socket(socketFile.getPath(), UnixSocketAddressWithTransport.TransportType.UDS_STREAM)
             .port(0)
             .queueSize(1)
             .timeout(500)  // non-zero timeout to ensure exception triggered if socket buffer full.
@@ -71,7 +64,7 @@ public SocketAddress call() throws Exception {
             .build();
 
         clientAggregate = new NonBlockingStatsDClientBuilder().prefix("my.prefix")
-            .addressLookup(addressLookup)
+            .socket(socketFile.getPath(), UnixSocketAddressWithTransport.TransportType.UDS_STREAM)
             .port(0)
             .queueSize(1)
             .timeout(500)  // non-zero timeout to ensure exception triggered if socket buffer full.