More robust implementation of a stream server

Author: iksaif

src/test/java/com/timgroup/statsd/DummyStatsDServer.java

@@ -35,12 +35,7 @@ public void run() {
                             ((Buffer)packet).clear();  // Cast necessary to handle Java9 covariant return types
                                                        // see: https://jira.mongodb.org/browse/JAVA-2559 for ref.
                             receive(packet);
-                            packetsReceived.addAndGet(1);
-
-                            packet.flip();
-                            for (String msg : StandardCharsets.UTF_8.decode(packet).toString().split("\n")) {
-                                addMessage(msg);
-                            }
+                            handlePacket(packet);
                         } catch (IOException e) {
                         }
                     }
@@ -51,6 +46,25 @@ public void run() {
         thread.start();
     }
 
+    protected boolean sleepIfFrozen() {
+        if (freeze) {
+            try {
+                Thread.sleep(10);
+            } catch (InterruptedException e) {
+            }
+        }
+        return freeze;
+    }
+
+    protected void handlePacket(ByteBuffer packet) {
+        packetsReceived.addAndGet(1);
+
+        packet.flip();
+        for (String msg : StandardCharsets.UTF_8.decode(packet).toString().split("\n")) {
+            addMessage(msg);
+        }
+    }
+
     public void waitForMessage() {
         waitForMessage(null);
     }

src/test/java/com/timgroup/statsd/UnixDatagramSocketDummyStatsDServer.java

@@ -1,11 +1,15 @@
 package com.timgroup.statsd;
 
 import java.io.IOException;
+import java.nio.Buffer;
 import java.nio.ByteBuffer;
 import java.nio.channels.DatagramChannel;
+import java.nio.charset.StandardCharsets;
 import jnr.unixsocket.UnixDatagramChannel;
 import jnr.unixsocket.UnixSocketAddress;
 
+import static com.timgroup.statsd.NonBlockingStatsDClient.DEFAULT_UDS_MAX_PACKET_SIZE_BYTES;
+
 public class UnixDatagramSocketDummyStatsDServer extends DummyStatsDServer {
     private final DatagramChannel server;
 
@@ -14,12 +18,12 @@ public UnixDatagramSocketDummyStatsDServer(String socketPath) throws IOException
         server.bind(new UnixSocketAddress(socketPath));
         this.listen();
     }
+
     @Override
     protected boolean isOpen() {
         return server.isOpen();
     }
 
-    @Override
     protected void receive(ByteBuffer packet) throws IOException {
         server.receive(packet);
     }

src/test/java/com/timgroup/statsd/UnixStreamSocketDummyStatsDServer.java

@@ -1,5 +1,6 @@
 package com.timgroup.statsd;
 
+import com.sun.xml.internal.ws.api.message.Packet;
 import java.io.IOException;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
@@ -22,14 +23,14 @@
 public class UnixStreamSocketDummyStatsDServer extends DummyStatsDServer {
     private final UnixServerSocketChannel server;
     private final ConcurrentLinkedQueue<UnixSocketChannel> channels = new ConcurrentLinkedQueue<>();
+    private final ConcurrentLinkedQueue<Packet> packets = new ConcurrentLinkedQueue<>();
 
     private final Logger logger = Logger.getLogger(UnixStreamSocketDummyStatsDServer.class.getName());
 
     public UnixStreamSocketDummyStatsDServer(String socketPath) throws IOException {
         server = UnixServerSocketChannel.open();
-        server.configureBlocking(false);
+        server.configureBlocking(true);
         server.socket().bind(new UnixSocketAddress(socketPath));
-        this.accept();
         this.listen();
     }
 
@@ -38,20 +39,23 @@ protected boolean isOpen() {
         return server.isOpen();
     }
 
-    protected void accept() {
+    @Override
+    protected void receive(ByteBuffer packet) throws IOException {
+        // This is unused because we re-implement listen() to fit our needs
+    }
+
+    @Override
+    protected void listen() {
+        logger.info("Listening on " + server.getLocalSocketAddress());
         Thread thread = new Thread(new Runnable() {
             @Override
             public void run() {
-                final ByteBuffer packet = ByteBuffer.allocate(DEFAULT_UDS_MAX_PACKET_SIZE_BYTES);
-
                 while(isOpen()) {
-                    if (freeze) {
-                        try {
-                            Thread.sleep(10);
-                        } catch (InterruptedException e) {
-                        }
-                    } else {
+                    if (sleepIfFrozen()) {
+                        continue;
+                    }
                         try {
+                            logger.info("Waiting for connection");
                             UnixSocketChannel clientChannel = server.accept();
                             if (clientChannel != null) {
                                 clientChannel.configureBlocking(true);
@@ -61,33 +65,54 @@ public void run() {
                                     logger.warning("Failed to get remote socket address");
                                 }
                                 channels.add(clientChannel);
+                                readChannel(clientChannel);
                             }
                         } catch (IOException e) {
                         }
-                    }
                 }
             }
         });
         thread.setDaemon(true);
         thread.start();
     }
 
-    @Override
-    protected void receive(ByteBuffer packet) throws IOException {
-        for (UnixSocketChannel channel : channels) {
-            if (channel.isConnected()) {
-                if (readPacket(channel, packet)) {
-                    return;
+    public void readChannel(final UnixSocketChannel clientChannel) {
+            logger.info("Reading from " + clientChannel);
+            Thread thread = new Thread(new Runnable() {
+            @Override
+            public void run() {
+                final ByteBuffer packet = ByteBuffer.allocate(DEFAULT_UDS_MAX_PACKET_SIZE_BYTES);
+
+                while(clientChannel.isOpen()) {
+                    if (sleepIfFrozen()) {
+                        continue;
+                    }
+                    ((Buffer)packet).clear();  // Cast necessary to handle Java9 covariant return types
+                    // see: https://jira.mongodb.org/browse/JAVA-2559 for ref.
+                    if (readPacket(clientChannel, packet)) {
+                        handlePacket(packet);
+                    } else {
+                        try {
+                            clientChannel.close();
+                        } catch (IOException e) {
+                            logger.warning("Failed to close channel: " + e);
+                        }
+                    }
+
                 }
+                logger.info("Disconnected from " + clientChannel);
             }
-        }
+        });
+        thread.setDaemon(true);
+        thread.start();
     }
 
-    private boolean readPacket(SocketChannel channel, ByteBuffer packet) throws IOException {
+    private boolean readPacket(SocketChannel channel, ByteBuffer packet) {
         try {
             ByteBuffer delimiterBuffer = ByteBuffer.allocate(Integer.SIZE / Byte.SIZE).order(ByteOrder.LITTLE_ENDIAN);
 
             int read = channel.read(delimiterBuffer);
+
             delimiterBuffer.flip();
             if (read <= 0) {
                 // There was nothing to read
@@ -100,16 +125,11 @@ private boolean readPacket(SocketChannel channel, ByteBuffer packet) throws IOEx
             }
 
             packet.limit(packetSize);
-            long deadline = System.nanoTime() + 1_000L * 1_000_000L;
-            while (packet.hasRemaining()) {
-                read = channel.read(packet);
-                if (deadline < System.nanoTime()) {
-                    channel.close();
-                }
+            while (packet.hasRemaining() && channel.isConnected()) {
+                channel.read(packet);
             }
             return true;
         } catch (IOException e) {
-            channel.close();
             return false;
         }
     }

src/test/java/com/timgroup/statsd/UnixStreamSocketTest.java

@@ -62,7 +62,7 @@ public SocketAddress call() throws Exception {
             .addressLookup(addressLookup)
             .port(0)
             .queueSize(1)
-            .timeout(1)  // non-zero timeout to ensure exception triggered if socket buffer full.
+            .timeout(500)  // non-zero timeout to ensure exception triggered if socket buffer full.
             .connectionTimeout(500)
             .socketBufferSize(1024 * 1024)
             .enableAggregation(false)
@@ -74,7 +74,7 @@ public SocketAddress call() throws Exception {
             .addressLookup(addressLookup)
             .port(0)
             .queueSize(1)
-            .timeout(1)  // non-zero timeout to ensure exception triggered if socket buffer full.
+            .timeout(500)  // non-zero timeout to ensure exception triggered if socket buffer full.
             .connectionTimeout(500)
             .socketBufferSize(1024 * 1024)
             .enableAggregation(false)
@@ -164,7 +164,7 @@ public void resist_dsd_timeout() throws Exception {
 
         while (lastException.getMessage() == null) {
             client.gauge("mycount", 20);
-            Thread.sleep(10);  // We need to fill the buffer, setting a shorter sleep
+
         }
         String excMessage = "Write timed out";
         assertThat(lastException.getMessage(), containsString(excMessage));