Update getInputStream to use a selector.

Author: sarahchen6

utils/socket-utils/src/main/java17/datadog/common/socket/TunnelingJdkSocket.java

@@ -9,7 +9,10 @@
 import java.net.SocketAddress;
 import java.net.SocketException;
 import java.net.UnixDomainSocketAddress;
+import java.nio.ByteBuffer;
 import java.nio.channels.Channels;
+import java.nio.channels.SelectionKey;
+import java.nio.channels.Selector;
 import java.nio.channels.SocketChannel;
 import java.nio.file.Path;
 
@@ -118,7 +121,46 @@ public InputStream getInputStream() throws IOException {
     if (isInputShutdown()) {
       throw new SocketException("Socket input is shutdown");
     }
-    return Channels.newInputStream(unixSocketChannel);
+
+    Selector selector = Selector.open();
+    unixSocketChannel.configureBlocking(false);
+    unixSocketChannel.register(selector, SelectionKey.OP_READ);
+    ByteBuffer buffer = ByteBuffer.allocate(256); // arbitrary buffer size for now
+
+    try {
+      if (selector.select(timeout) == 0) {
+        System.out.println("Timeout (" + timeout + "ms) while waiting for data.");
+      }
+      for (SelectionKey key : selector.selectedKeys()) {
+        if (key.isReadable()) {
+          int r = unixSocketChannel.read(buffer);
+          if (r == -1) {
+            unixSocketChannel.close();
+            System.out.println("Not accepting client messages anymore.");
+          }
+        }
+      }
+      buffer.flip();
+    } finally {
+      selector.close();
+    }
+
+    return new InputStream() {
+      @Override
+      public int read() {
+        return buffer.hasRemaining() ? (buffer.get() & 0xFF) : -1;
+      }
+
+      @Override
+      public int read(byte[] bytes, int off, int len) {
+        if (!buffer.hasRemaining()) {
+          return -1;
+        }
+        len = Math.min(len, buffer.remaining());
+        buffer.get(bytes, off, len);
+        return len;
+      }
+    };
   }
 
   @Override

utils/socket-utils/src/test/java/datadog/common/socket/TunnelingJdkSocketTest.java

@@ -31,7 +31,7 @@ public void testTimeout() throws Exception {
     // create client
     TunnelingJdkSocket clientSocket = createClient(socketPath);
 
-    // expect a failure after three seconds because timeout is not supported yet
+    // will fail after three seconds if timeout (set to one second) is not supported
     assertTimeoutPreemptively(Duration.ofMillis(3000), () -> clientSocket.getInputStream().read());
 
     // clean up