Adding capability to support Immutable clients by the Worker and Worker.Builder

src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/Worker.java

@@ -32,17 +32,24 @@ import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
+import org.apache.commons.lang.StringUtils;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
-import com.amazonaws.regions.Region;
+import com.amazonaws.AmazonWebServiceClient;
+import com.amazonaws.ClientConfiguration;
+import com.amazonaws.auth.AWSCredentialsProvider;
+import com.amazonaws.client.builder.AwsClientBuilder;
 import com.amazonaws.regions.RegionUtils;
 import com.amazonaws.services.cloudwatch.AmazonCloudWatch;
 import com.amazonaws.services.cloudwatch.AmazonCloudWatchClient;
+import com.amazonaws.services.cloudwatch.AmazonCloudWatchClientBuilder;
 import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
 import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient;
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder;
 import com.amazonaws.services.kinesis.AmazonKinesis;
 import com.amazonaws.services.kinesis.AmazonKinesisClient;
+import com.amazonaws.services.kinesis.AmazonKinesisClientBuilder;
 import com.amazonaws.services.kinesis.clientlibrary.interfaces.ICheckpoint;
 import com.amazonaws.services.kinesis.clientlibrary.interfaces.v2.IRecordProcessor;
 import com.amazonaws.services.kinesis.clientlibrary.interfaces.v2.IRecordProcessorFactory;
@@ -60,6 +67,7 @@ import com.amazonaws.services.kinesis.metrics.interfaces.MetricsLevel;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 
+import lombok.NonNull;
 import lombok.Setter;
 import lombok.experimental.Accessors;
 
@@ -396,28 +404,16 @@ public class Worker implements Runnable {
 
         // If a region name was explicitly specified, use it as the region for Amazon Kinesis and Amazon DynamoDB.
         if (config.getRegionName() != null) {
-            Region region = RegionUtils.getRegion(config.getRegionName());
+            setRegionForClient((AmazonWebServiceClient) kinesisClient, config.getRegionName());
-            kinesisClient.setRegion(region);
+            setRegionForClient((AmazonWebServiceClient) dynamoDBClient, config.getRegionName());
-            LOG.debug("The region of Amazon Kinesis client has been set to " + config.getRegionName());
-            dynamoDBClient.setRegion(region);
-            LOG.debug("The region of Amazon DynamoDB client has been set to " + config.getRegionName());
         }
         // If a dynamoDB endpoint was explicitly specified, use it to set the DynamoDB endpoint.
         if (config.getDynamoDBEndpoint() != null) {
-            dynamoDBClient.setEndpoint(config.getDynamoDBEndpoint());
+            setEndpointForClient((AmazonWebServiceClient) dynamoDBClient, config.getDynamoDBEndpoint());
-            LOG.debug("The endpoint of Amazon DynamoDB client has been set to " + config.getDynamoDBEndpoint());
         }
         // If a kinesis endpoint was explicitly specified, use it to set the region of kinesis.
         if (config.getKinesisEndpoint() != null) {
-            kinesisClient.setEndpoint(config.getKinesisEndpoint());
+            setEndpointForClient((AmazonWebServiceClient) kinesisClient, config.getKinesisEndpoint());
-            if (config.getRegionName() != null) {
-                LOG.warn("Received configuration for both region name as " + config.getRegionName()
-                        + ", and Amazon Kinesis endpoint as " + config.getKinesisEndpoint()
-                        + ". Amazon Kinesis endpoint will overwrite region name.");
-                LOG.debug("The region of Amazon Kinesis client has been overwritten to " + config.getKinesisEndpoint());
-            } else {
-                LOG.debug("The region of Amazon Kinesis client has been set to " + config.getKinesisEndpoint());
-            }
         }
     }
 
@@ -1068,9 +1064,7 @@ public class Worker implements Runnable {
             metricsFactory = new NullMetricsFactory();
         } else {
             if (config.getRegionName() != null) {
-                Region region = RegionUtils.getRegion(config.getRegionName());
+                setRegionForClient((AmazonWebServiceClient) cloudWatchClient, config.getRegionName());
-                cloudWatchClient.setRegion(region);
-                LOG.debug("The region of Amazon CloudWatch client has been set to " + config.getRegionName());
             }
             metricsFactory = new WorkerCWMetricsFactory(cloudWatchClient, config.getApplicationName(),
                     config.getMetricsBufferTimeMillis(), config.getMetricsMaxQueueSize(), config.getMetricsLevel(),
@@ -1089,6 +1083,24 @@ public class Worker implements Runnable {
         return new WorkerThreadPoolExecutor(threadFactory);
     }
 
+    private static void setRegionForClient(@NonNull final AmazonWebServiceClient client, @NonNull final String region) {
+        try {
+            client.setRegion(RegionUtils.getRegion(region));
+            LOG.debug("Region set for client to " + region);
+        } catch (UnsupportedOperationException e) {
+            LOG.debug("Trying to set region for immutable client", e);
+        }
+    }
+
+    private static void setEndpointForClient(@NonNull final AmazonWebServiceClient client, @NonNull final String endpointUrl) {
+        try {
+            client.setEndpoint(endpointUrl);
+            LOG.debug("Endpoint set for client to " + endpointUrl);
+        } catch (UnsupportedOperationException e) {
+            LOG.debug("Trying to set endpoint for immutable client", e);
+        }
+    }
+
     /**
      * Extension to CWMetricsFactory, so worker can identify whether it owns the metrics factory instance or not.
      * Visible and non-final only for testing.
@@ -1140,6 +1152,21 @@ public class Worker implements Runnable {
         @Setter @Accessors(fluent = true)
         private IKinesisProxy kinesisProxy;
 
+        @VisibleForTesting
+        AmazonKinesis getKinesisClient() {
+            return kinesisClient;
+        }
+
+        @VisibleForTesting
+        AmazonDynamoDB getDynamoDBClient() {
+            return dynamoDBClient;
+        }
+
+        @VisibleForTesting
+        AmazonCloudWatch getCloudWatchClient() {
+            return cloudWatchClient;
+        }
+
         /**
          * Provide a V1 {@link com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessor
          * IRecordProcessor}.
@@ -1187,44 +1214,39 @@ public class Worker implements Runnable {
                 execService = getExecutorService();
             }
             if (kinesisClient == null) {
-                kinesisClient = new AmazonKinesisClient(config.getKinesisCredentialsProvider(),
+                kinesisClient = (AmazonKinesis) createClient(AmazonKinesisClientBuilder.standard(),
-                        config.getKinesisClientConfiguration());
+                        config.getKinesisCredentialsProvider(),
+                        config.getKinesisClientConfiguration(),
+                        config.getKinesisEndpoint(),
+                        config.getRegionName());
             }
             if (dynamoDBClient == null) {
-                dynamoDBClient = new AmazonDynamoDBClient(config.getDynamoDBCredentialsProvider(),
+                dynamoDBClient = (AmazonDynamoDB) createClient(AmazonDynamoDBClientBuilder.standard(),
-                        config.getDynamoDBClientConfiguration());
+                        config.getDynamoDBCredentialsProvider(),
+                        config.getDynamoDBClientConfiguration(),
+                        config.getDynamoDBEndpoint(),
+                        config.getRegionName());
             }
             if (cloudWatchClient == null) {
-                cloudWatchClient = new AmazonCloudWatchClient(config.getCloudWatchCredentialsProvider(),
+                cloudWatchClient = (AmazonCloudWatch) createClient(AmazonCloudWatchClientBuilder.standard(),
-                        config.getCloudWatchClientConfiguration());
+                        config.getCloudWatchCredentialsProvider(),
+                        config.getCloudWatchClientConfiguration(),
+                        null,
+                        config.getRegionName());
             }
             // If a region name was explicitly specified, use it as the region for Amazon Kinesis and Amazon DynamoDB.
             if (config.getRegionName() != null) {
-                Region region = RegionUtils.getRegion(config.getRegionName());
+                setRegionForClient((AmazonWebServiceClient) cloudWatchClient, config.getRegionName());
-                cloudWatchClient.setRegion(region);
+                setRegionForClient((AmazonWebServiceClient) kinesisClient, config.getRegionName());
-                LOG.debug("The region of Amazon CloudWatch client has been set to " + config.getRegionName());
+                setRegionForClient((AmazonWebServiceClient) dynamoDBClient, config.getRegionName());
-                kinesisClient.setRegion(region);
-                LOG.debug("The region of Amazon Kinesis client has been set to " + config.getRegionName());
-                dynamoDBClient.setRegion(region);
-                LOG.debug("The region of Amazon DynamoDB client has been set to " + config.getRegionName());
             }
             // If a dynamoDB endpoint was explicitly specified, use it to set the DynamoDB endpoint.
             if (config.getDynamoDBEndpoint() != null) {
-                dynamoDBClient.setEndpoint(config.getDynamoDBEndpoint());
+                setEndpointForClient((AmazonWebServiceClient) dynamoDBClient, config.getDynamoDBEndpoint());
-                LOG.debug("The endpoint of Amazon DynamoDB client has been set to " + config.getDynamoDBEndpoint());
             }
             // If a kinesis endpoint was explicitly specified, use it to set the region of kinesis.
             if (config.getKinesisEndpoint() != null) {
-                kinesisClient.setEndpoint(config.getKinesisEndpoint());
+                setEndpointForClient((AmazonWebServiceClient) kinesisClient, config.getKinesisEndpoint());
-                if (config.getRegionName() != null) {
-                    LOG.warn("Received configuration for both region name as " + config.getRegionName()
-                            + ", and Amazon Kinesis endpoint as " + config.getKinesisEndpoint()
-                            + ". Amazon Kinesis endpoint will overwrite region name.");
-                    LOG.debug("The region of Amazon Kinesis client has been overwritten to "
-                            + config.getKinesisEndpoint());
-                } else {
-                    LOG.debug("The region of Amazon Kinesis client has been set to " + config.getKinesisEndpoint());
-                }
             }
             if (metricsFactory == null) {
                 metricsFactory = getMetricsFactory(cloudWatchClient, config);
@@ -1274,5 +1296,27 @@ public class Worker implements Runnable {
 
         }
 
+        AmazonWebServiceClient createClient(@NonNull final AwsClientBuilder builder,
+                                            final AWSCredentialsProvider credentialsProvider,
+                                            final ClientConfiguration clientConfiguration,
+                                            final String endpointUrl,
+                                            final String region) {
+            if (credentialsProvider != null) {
+                builder.withCredentials(credentialsProvider);
+            }
+            if (clientConfiguration != null) {
+                builder.withClientConfiguration(clientConfiguration);
+            }
+            if (StringUtils.isNotEmpty(endpointUrl)) {
+                LOG.warn("Received configuration for both region name as " + region + ", and endpoint as "
+                        + endpointUrl + ". Client endpoint will overwrite region name.");
+                builder.withEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration(endpointUrl, region));
+            } else if (StringUtils.isNotEmpty(region)) {
+                LOG.debug("The region for the client has been set to " + region);
+                builder.withRegion(region);
+            }
+            return (AmazonWebServiceClient) builder.build();
+        }
+
     }
 }

src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/KinesisClientLibConfigurationTest.java

@@ -268,8 +268,8 @@ public class KinesisClientLibConfigurationTest {
         IRecordProcessorFactory processorFactory = Mockito.mock(IRecordProcessorFactory.class);
         new Worker(processorFactory, kclConfig);
 
-        Mockito.verify(kclConfig, Mockito.times(9)).getRegionName();
+        Mockito.verify(kclConfig, Mockito.times(5)).getRegionName();
-        Mockito.verify(kclConfig, Mockito.times(4)).getKinesisEndpoint();
+        Mockito.verify(kclConfig, Mockito.times(2)).getKinesisEndpoint();
 
         kclConfig = Mockito.spy(
                 new KinesisClientLibConfiguration("Test", "Test", credentialsProvider, "0")
@@ -277,8 +277,8 @@ public class KinesisClientLibConfigurationTest {
 
         new Worker(processorFactory, kclConfig);
 
-        Mockito.verify(kclConfig, Mockito.times(3)).getRegionName();
+        Mockito.verify(kclConfig, Mockito.times(2)).getRegionName();
-        Mockito.verify(kclConfig, Mockito.times(3)).getKinesisEndpoint();
+        Mockito.verify(kclConfig, Mockito.times(2)).getKinesisEndpoint();
     }
 
     @Test

src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/WorkerTest.java

@@ -19,6 +19,9 @@ import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.CoreMatchers.isA;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
@@ -63,6 +66,8 @@ import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 
+import com.amazonaws.ClientConfiguration;
+import com.amazonaws.client.builder.AwsClientBuilder;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.hamcrest.Condition;
@@ -80,8 +85,14 @@ import org.mockito.invocation.InvocationOnMock;
 import org.mockito.runners.MockitoJUnitRunner;
 import org.mockito.stubbing.Answer;
 
+import com.amazonaws.regions.RegionUtils;
+import com.amazonaws.services.cloudwatch.AmazonCloudWatch;
+import com.amazonaws.services.cloudwatch.AmazonCloudWatchClientBuilder;
 import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder;
 import com.amazonaws.services.dynamodbv2.local.embedded.DynamoDBEmbedded;
+import com.amazonaws.services.kinesis.AmazonKinesis;
+import com.amazonaws.services.kinesis.AmazonKinesisClientBuilder;
 import com.amazonaws.services.kinesis.clientlibrary.exceptions.KinesisClientLibNonRetryableException;
 import com.amazonaws.services.kinesis.clientlibrary.interfaces.ICheckpoint;
 import com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessorCheckpointer;
@@ -90,8 +101,8 @@ import com.amazonaws.services.kinesis.clientlibrary.interfaces.v2.IRecordProcess
 import com.amazonaws.services.kinesis.clientlibrary.lib.worker.Worker.WorkerCWMetricsFactory;
 import com.amazonaws.services.kinesis.clientlibrary.lib.worker.Worker.WorkerThreadPoolExecutor;
 import com.amazonaws.services.kinesis.clientlibrary.proxies.IKinesisProxy;
-import com.amazonaws.services.kinesis.clientlibrary.proxies.KinesisProxy;
 import com.amazonaws.services.kinesis.clientlibrary.proxies.KinesisLocalFileProxy;
+import com.amazonaws.services.kinesis.clientlibrary.proxies.KinesisProxy;
 import com.amazonaws.services.kinesis.clientlibrary.proxies.util.KinesisLocalFileDataCreator;
 import com.amazonaws.services.kinesis.clientlibrary.types.ExtendedSequenceNumber;
 import com.amazonaws.services.kinesis.clientlibrary.types.InitializationInput;
@@ -1527,6 +1538,115 @@ public class WorkerTest {
         Assert.assertSame(leaseManager, worker.getLeaseCoordinator().getLeaseManager());
     }
 
+    @Test
+    public void testBuilderSetRegionAndEndpointToClient() {
+        IRecordProcessorFactory recordProcessorFactory = mock(IRecordProcessorFactory.class);
+        final String endpoint = "TestEndpoint";
+        KinesisClientLibConfiguration config = new KinesisClientLibConfiguration("TestApp", null, null, null)
+                .withRegionName("us-west-2")
+                .withKinesisEndpoint(endpoint)
+                .withDynamoDBEndpoint(endpoint);
+
+        AmazonKinesis kinesisClient = spy(AmazonKinesisClientBuilder.defaultClient());
+        AmazonDynamoDB dynamoDBClient = spy(AmazonDynamoDBClientBuilder.defaultClient());
+        AmazonCloudWatch cloudWatchClient = spy(AmazonCloudWatchClientBuilder.defaultClient());
+
+        new Worker.Builder().recordProcessorFactory(recordProcessorFactory).config(config)
+                .kinesisClient(kinesisClient)
+                .dynamoDBClient(dynamoDBClient)
+                .cloudWatchClient(cloudWatchClient)
+                .build();
+
+        verify(kinesisClient, times(1)).setRegion(eq(RegionUtils.getRegion(config.getRegionName())));
+        verify(dynamoDBClient, times(1)).setRegion(eq(RegionUtils.getRegion(config.getRegionName())));
+        verify(cloudWatchClient, times(2)).setRegion(eq(RegionUtils.getRegion(config.getRegionName())));
+
+        verify(kinesisClient, times(1)).setEndpoint(eq(endpoint));
+        verify(dynamoDBClient, times(1)).setEndpoint(eq(endpoint));
+        verify(cloudWatchClient, never()).setEndpoint(anyString());
+    }
+
+    @Test
+    public void testBuilderSetRegionToClient() {
+        IRecordProcessorFactory recordProcessorFactory = mock(IRecordProcessorFactory.class);
+        String region = "us-west-2";
+        KinesisClientLibConfiguration config = new KinesisClientLibConfiguration("TestApp", null, null, null)
+                .withRegionName(region);
+
+        Worker.Builder builder = new Worker.Builder();
+
+        AmazonKinesis kinesisClient = spy(AmazonKinesisClientBuilder.defaultClient());
+        AmazonDynamoDB dynamoDBClient = spy(AmazonDynamoDBClientBuilder.defaultClient());
+        AmazonCloudWatch cloudWatchClient = spy(AmazonCloudWatchClientBuilder.defaultClient());
+
+        builder.recordProcessorFactory(recordProcessorFactory).config(config)
+                .kinesisClient(kinesisClient)
+                .dynamoDBClient(dynamoDBClient)
+                .cloudWatchClient(cloudWatchClient)
+                .build();
+
+        verify(kinesisClient, times(1)).setRegion(eq(RegionUtils.getRegion(config.getRegionName())));
+        verify(dynamoDBClient, times(1)).setRegion(eq(RegionUtils.getRegion(config.getRegionName())));
+        verify(cloudWatchClient, times(2)).setRegion(eq(RegionUtils.getRegion(config.getRegionName())));
+
+        verify(kinesisClient, never()).setEndpoint(any());
+        verify(dynamoDBClient, never()).setEndpoint(any());
+        verify(cloudWatchClient, never()).setEndpoint(any());
+    }
+
+    @Test
+    public void testBuilderGenerateClients() {
+        IRecordProcessorFactory recordProcessorFactory = mock(IRecordProcessorFactory.class);
+        KinesisClientLibConfiguration config = new KinesisClientLibConfiguration("TestApp", null, null, null);
+        Worker.Builder builder = spy(new Worker.Builder().recordProcessorFactory(recordProcessorFactory).config(config));
+
+        assertNull(builder.getKinesisClient());
+        assertNull(builder.getDynamoDBClient());
+        assertNull(builder.getCloudWatchClient());
+
+        builder.build();
+
+        assertTrue(builder.getKinesisClient() instanceof AmazonKinesis);
+        assertTrue(builder.getDynamoDBClient() instanceof AmazonDynamoDB);
+        assertTrue(builder.getCloudWatchClient() instanceof AmazonCloudWatch);
+
+        verify(builder, times(3)).createClient(
+                any(AwsClientBuilder.class), eq(null), any(ClientConfiguration.class), eq(null), eq(null));
+    }
+
+    @Test
+    public void testBuilderGenerateClientsWithRegion() {
+        IRecordProcessorFactory recordProcessorFactory = mock(IRecordProcessorFactory.class);
+        String region = "us-west-2";
+        KinesisClientLibConfiguration config = new KinesisClientLibConfiguration("TestApp", null, null, null)
+                .withRegionName(region);
+
+        Worker.Builder builder = spy(new Worker.Builder());
+
+        builder.recordProcessorFactory(recordProcessorFactory).config(config).build();
+
+        verify(builder, times(3)).createClient(
+                any(AwsClientBuilder.class), eq(null), any(ClientConfiguration.class), eq(null), eq(region));
+    }
+
+    @Test
+    public void testBuilderGenerateClientsWithEndpoint() {
+        IRecordProcessorFactory recordProcessorFactory = mock(IRecordProcessorFactory.class);
+        String region = "us-west-2";
+        String endpointUrl = "TestEndpoint";
+        KinesisClientLibConfiguration config = new KinesisClientLibConfiguration("TestApp", null, null, null)
+                .withRegionName(region).withKinesisEndpoint(endpointUrl).withDynamoDBEndpoint(endpointUrl);
+
+        Worker.Builder builder = spy(new Worker.Builder());
+
+        builder.recordProcessorFactory(recordProcessorFactory).config(config).build();
+
+        verify(builder, times(2)).createClient(
+                any(AwsClientBuilder.class), eq(null), any(ClientConfiguration.class), eq(endpointUrl), eq(region));
+        verify(builder, times(1)).createClient(
+                any(AwsClientBuilder.class), eq(null), any(ClientConfiguration.class), eq(null), eq(region));
+    }
+
     private abstract class InjectableWorker extends Worker {
         InjectableWorker(String applicationName, IRecordProcessorFactory recordProcessorFactory,
                 KinesisClientLibConfiguration config, StreamConfig streamConfig,