add socket eos detection

Author: Mark Kuhn

src/main/java/software/amazon/cloudwatchlogs/emf/sinks/TCPClient.java

@@ -20,6 +20,8 @@
 import java.io.OutputStream;
 import java.net.InetSocketAddress;
 import java.net.Socket;
+import java.util.concurrent.*;
+
 import lombok.extern.slf4j.Slf4j;
 
 /** A client that would connect to a TCP socket. */
@@ -51,9 +53,7 @@ protected Socket createSocket() {
 
     @Override
     public synchronized void sendMessage(String message) {
-        if (socket == null || socket.isClosed() || shouldConnect) {
-            connect();
-        }
+        checkConnection();
 
         OutputStream os;
         try {
@@ -72,6 +72,48 @@ public synchronized void sendMessage(String message) {
         }
     }
 
+    /**
+     * Performs checks to see if the socket is connected.
+     * If not, it will attempt to connect.
+     */
+    private void checkConnection() {
+        if (socket == null || socket.isClosed() || shouldConnect) {
+            connect();
+            return;
+        }
+
+        try {
+            if (streamIsClosed()) {
+                connect();
+            }
+        } catch (RuntimeException | IOException e) {
+            connect();
+        }
+    }
+
+    /**
+     * Checks if the socket's input stream has reached EOS.
+     *
+     * @return true if the input stream is closed, false otherwise.
+     */
+    private boolean streamIsClosed() throws IOException {
+        Integer result;
+        Callable<Integer> readTask = socket.getInputStream()::read;
+
+        ExecutorService executor = Executors.newCachedThreadPool();
+        Future<Integer> readFuture = executor.submit(readTask);
+
+        try {
+            result = readFuture.get(1, TimeUnit.MILLISECONDS);
+        } catch (ExecutionException | InterruptedException e) {
+            throw new RuntimeException(e);
+        } catch (TimeoutException ignored) {
+            return false;
+        }
+
+        return result == -1;
+    }
+
     @Override
     public void close() throws IOException {
         if (socket != null) {

src/test/java/software/amazon/cloudwatchlogs/emf/sinks/TCPClientTest.java

@@ -22,6 +22,8 @@
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.InputStream;
+import java.net.ServerSocket;
 import java.net.Socket;
 import org.junit.Test;
 
@@ -48,4 +50,55 @@ protected Socket createSocket() {
 
         assertEquals(bos.toString(), message);
     }
+
+    @Test(timeout = 5000)
+    public void testSendMessageAfterAgentReconnect() throws InterruptedException {
+        String[] messages = new String[2];
+
+        Thread server = new Thread(() -> {
+            try {
+                ServerSocket serverSocket = new ServerSocket(Endpoint.DEFAULT_TCP_ENDPOINT.getPort());
+                Socket socket = serverSocket.accept();
+                InputStream is = socket.getInputStream();
+                byte[] bytes = new byte[1024];
+                int read = is.read(bytes);
+                messages[0] = new String(bytes, 0, read);
+
+                // Disconnect
+                socket.close();
+                serverSocket.close();
+
+                // Reconnect
+                serverSocket = new ServerSocket(Endpoint.DEFAULT_TCP_ENDPOINT.getPort());
+                socket = serverSocket.accept();
+                is = socket.getInputStream();
+                bytes = new byte[1024];
+                read = is.read(bytes);
+                messages[1] = new String(bytes, 0, read);
+
+                socket.close();
+                serverSocket.close();
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        });
+
+        Thread client = new Thread(() -> {
+            try (TCPClient tcpClient = new TCPClient(Endpoint.DEFAULT_TCP_ENDPOINT)) {
+                tcpClient.sendMessage("Test message 1");
+                Thread.sleep(1000);
+                tcpClient.sendMessage("Test message 2");
+            } catch (InterruptedException | IOException e) {
+                throw new RuntimeException(e);
+            }
+        });
+
+        server.start();
+        client.start();
+        server.join();
+        client.join();
+
+        assertEquals("Test message 1", messages[0]);
+        assertEquals("Test message 2", messages[1]);
+    }
 }