add retry queue

Author: Mark Kuhn

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

@@ -20,6 +20,7 @@
 import java.io.OutputStream;
 import java.net.InetSocketAddress;
 import java.net.Socket;
+import java.util.LinkedList;
 import lombok.extern.slf4j.Slf4j;
 
 /** A client that would connect to a TCP socket. */
@@ -29,6 +30,8 @@ public class TCPClient implements SocketClient {
     private final Endpoint endpoint;
     private Socket socket;
     private boolean shouldConnect = true;
+    private final LinkedList<String> retryQueue = new LinkedList<>();
+    private String lastMessageSent;
 
     public TCPClient(Endpoint endpoint) {
         this.endpoint = endpoint;
@@ -51,27 +54,45 @@ protected Socket createSocket() {
 
     @Override
     public synchronized void sendMessage(String message) {
-        if (socket == null || socket.isClosed() || shouldConnect) {
-            connect();
-        }
+        checkConnection();
 
         OutputStream os;
         try {
             os = socket.getOutputStream();
         } catch (IOException e) {
             shouldConnect = true;
+            retryQueue.add(message);
             throw new RuntimeException(
                     "Failed to write message to the socket. Failed to open output stream.", e);
         }
 
         try {
+            while (!retryQueue.isEmpty()) {
+                String retryMessage = retryQueue.peek();
+                os.write(retryMessage.getBytes());
+                retryQueue.pop();
+            }
+
             os.write(message.getBytes());
-        } catch (Exception e) {
+            lastMessageSent = message;
+        } catch (IOException e) {
+            // For broken pipe exception, retry last sent message
+            if (e.getMessage().contains("Broken pipe")) {
+                retryQueue.add(lastMessageSent);
+            }
+
             shouldConnect = true;
+            retryQueue.add(message);
             throw new RuntimeException("Failed to write message to the socket.", e);
         }
     }
 
+    private void checkConnection() {
+        if (socket == null || socket.isClosed() || shouldConnect) {
+            connect();
+        }
+    }
+
     @Override
     public void close() throws IOException {
         if (socket != null) {

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

@@ -22,6 +22,7 @@
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.net.ServerSocket;
 import java.net.Socket;
 import org.junit.Test;
 
@@ -48,4 +49,20 @@ protected Socket createSocket() {
 
         assertEquals(bos.toString(), message);
     }
+
+    @Test(timeout = 5000)
+    public void testSendMessageWithSocketServer() throws IOException {
+        TCPClient client = new TCPClient(new Endpoint("0.0.0.0", 9999, Protocol.TCP));
+        ServerSocket server = new ServerSocket(9999);
+        client.sendMessage("Test message");
+        Socket socket = server.accept();
+
+        byte[] bytes = new byte[1024];
+        int read = socket.getInputStream().read(bytes);
+        String message = new String(bytes, 0, read);
+        socket.close();
+        server.close();
+
+        assertEquals("Test message", message);
+    }
 }