Merge pull request #94 from awslabs/caowei

Fix TCP client seeing "already connected" error when resend message after having exception in the first send

README.md

@@ -324,11 +324,12 @@ We have 2 different types of tests:
 	```sh
 	export AWS_ACCESS_KEY_ID=YOUR_ACCESS_KEY_ID
 	export AWS_SECRET_ACCESS_KEY=YOUR_ACCESS_KEY
+	export AWS_SESSION_TOKEN=YOUR_AWS_SESSION_TOKEN
 	export AWS_REGION=us-west-2
 	./gradlew integ
 	```
 
-	**NOTE**: You need to replace the access key id and access key with your own AWS credentials.
+	**NOTE**: You need to replace the access key id and access key with your own AWS credentials. Another option is using IAM user access key pair without session token.
 
 ### Formatting
 

bin/start-agent.sh

@@ -6,6 +6,7 @@
 # usage:
 #   export AWS_ACCESS_KEY_ID=
 #   export AWS_SECRET_ACCESS_KEY=
+#   export AWS_SESSION_TOKEN=
 #   export AWS_REGION=us-west-2
 #   ./start-agent.sh
 
@@ -22,6 +23,7 @@ pushd $rootdir/src/integration-test/resources/agent
 echo "[AmazonCloudWatchAgent]
 aws_access_key_id = $AWS_ACCESS_KEY_ID
 aws_secret_access_key = $AWS_SECRET_ACCESS_KEY
+aws_session_token = $AWS_SESSION_TOKEN
 " > ./.aws/credentials
 
 echo "[profile AmazonCloudWatchAgent]
@@ -33,5 +35,6 @@ docker run  -p 25888:25888/udp -p 25888:25888/tcp  \
     -e AWS_REGION=$AWS_REGION \
     -e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \
     -e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \
+    -e AWS_SESSION_TOKEN=$AWS_SESSION_TOKEN \
     agent:latest &> $tempfile &
 popd

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

@@ -31,12 +31,12 @@ public class TCPClient implements SocketClient {
     private boolean shouldConnect = true;
 
     public TCPClient(Endpoint endpoint) {
-        socket = createSocket();
         this.endpoint = endpoint;
     }
 
     private void connect() {
         try {
+            socket = createSocket();
             socket.connect(new InetSocketAddress(endpoint.getHost(), endpoint.getPort()));
             shouldConnect = false;
         } catch (Exception e) {
@@ -51,7 +51,7 @@ protected Socket createSocket() {
 
     @Override
     public synchronized void sendMessage(String message) {
-        if (socket.isClosed() || shouldConnect) {
+        if (socket == null || socket.isClosed() || shouldConnect) {
             connect();
         }
 

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

@@ -48,4 +48,50 @@ protected Socket createSocket() {
 
         assertEquals(bos.toString(), message);
     }
+
+    @Test
+    public void testSendMessageWithGetOSException_THEN_createSocketTwice() throws IOException {
+        Socket socket = mock(Socket.class);
+        doNothing().when(socket).connect(any());
+        when(socket.getOutputStream()).thenThrow(IOException.class);
+
+        Endpoint endpoint = Endpoint.DEFAULT_TCP_ENDPOINT;
+        TCPClient client =
+                new TCPClient(endpoint) {
+                    @Override
+                    protected Socket createSocket() {
+                        return socket;
+                    }
+                };
+
+        TCPClient spyClient = spy(client);
+
+        String message = "Test message";
+        spyClient.sendMessage(message);
+        verify(spyClient, atLeast(2)).createSocket();
+    }
+
+    @Test
+    public void testSendMessageWithWriteOSException_THEN_createSocketTwice() throws IOException {
+        Socket socket = mock(Socket.class);
+        doNothing().when(socket).connect(any());
+        ByteArrayOutputStream bos = mock(ByteArrayOutputStream.class);
+        when(socket.getOutputStream()).thenReturn(bos);
+        doThrow(IOException.class).when(bos).write(any());
+
+        Endpoint endpoint = Endpoint.DEFAULT_TCP_ENDPOINT;
+        TCPClient client =
+                new TCPClient(endpoint) {
+                    @Override
+                    protected Socket createSocket() {
+                        return socket;
+                    }
+                };
+
+        TCPClient spyClient = spy(client);
+
+        String message = "Test message";
+        spyClient.sendMessage(message);
+        verify(spyClient, atLeast(2)).createSocket();
+    }
 }