TCP Client issue fix

Author: caowei-amazon

README.md

@@ -325,6 +325,7 @@ We have 2 different types of tests:
 	export AWS_ACCESS_KEY_ID=YOUR_ACCESS_KEY_ID
 	export AWS_SECRET_ACCESS_KEY=YOUR_ACCESS_KEY
 	export AWS_REGION=us-west-2
+ 	export AWS_SESSION_TOKEN=Your_AWS_SESSION_TOKEN
 	./gradlew integ
 	```
 

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

build.gradle

@@ -64,6 +64,7 @@ dependencies {
 	implementation 'com.fasterxml.jackson.core:jackson-annotations:2.11.1'
 	implementation 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310:2.11.1'
 	implementation 'org.slf4j:slf4j-api:1.7.30'
+	implementation 'org.slf4j:slf4j-simple:1.7.30'
 
 	// Use JUnit test framework
 	testImplementation 'software.amazon.awssdk:cloudwatch:2.13.54'

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

@@ -31,12 +31,13 @@ public class TCPClient implements SocketClient {
     private boolean shouldConnect = true;
 
     public TCPClient(Endpoint endpoint) {
-        socket = createSocket();
         this.endpoint = endpoint;
     }
 
     private void connect() {
         try {
+            // Avoid "socket already connected" error (https://issues.amazon.com/issues/P54323886)
+            socket = createSocket();
             socket.connect(new InetSocketAddress(endpoint.getHost(), endpoint.getPort()));
             shouldConnect = false;
         } catch (Exception e) {
@@ -51,7 +52,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();
+    }
 }