From 053bc226d0b9c9c7f0a1d73464f41d77073ef8d4 Mon Sep 17 00:00:00 2001 From: Sahil Palvia Date: Wed, 14 Feb 2018 10:29:05 -0800 Subject: [PATCH] Adding capability to support Immutable clients by the Worker and Worker.Builder --- .../clientlibrary/lib/worker/Worker.java | 134 ++++++++++++------ .../KinesisClientLibConfigurationTest.java | 8 +- .../clientlibrary/lib/worker/WorkerTest.java | 122 +++++++++++++++- 3 files changed, 214 insertions(+), 50 deletions(-) diff --git a/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/Worker.java b/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/Worker.java index 5b004cc5..c453bd34 100644 --- a/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/Worker.java +++ b/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()); - 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()); + setRegionForClient((AmazonWebServiceClient) kinesisClient, config.getRegionName()); + setRegionForClient((AmazonWebServiceClient) dynamoDBClient, config.getRegionName()); } // If a dynamoDB endpoint was explicitly specified, use it to set the DynamoDB endpoint. if (config.getDynamoDBEndpoint() != null) { - dynamoDBClient.setEndpoint(config.getDynamoDBEndpoint()); - LOG.debug("The endpoint of Amazon DynamoDB client has been set to " + config.getDynamoDBEndpoint()); + setEndpointForClient((AmazonWebServiceClient) dynamoDBClient, 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()); - 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()); - } + setEndpointForClient((AmazonWebServiceClient) kinesisClient, 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()); - cloudWatchClient.setRegion(region); - LOG.debug("The region of Amazon CloudWatch client has been set to " + config.getRegionName()); + setRegionForClient((AmazonWebServiceClient) cloudWatchClient, 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(), - config.getKinesisClientConfiguration()); + kinesisClient = (AmazonKinesis) createClient(AmazonKinesisClientBuilder.standard(), + config.getKinesisCredentialsProvider(), + config.getKinesisClientConfiguration(), + config.getKinesisEndpoint(), + config.getRegionName()); } if (dynamoDBClient == null) { - dynamoDBClient = new AmazonDynamoDBClient(config.getDynamoDBCredentialsProvider(), - config.getDynamoDBClientConfiguration()); + dynamoDBClient = (AmazonDynamoDB) createClient(AmazonDynamoDBClientBuilder.standard(), + config.getDynamoDBCredentialsProvider(), + config.getDynamoDBClientConfiguration(), + config.getDynamoDBEndpoint(), + config.getRegionName()); } if (cloudWatchClient == null) { - cloudWatchClient = new AmazonCloudWatchClient(config.getCloudWatchCredentialsProvider(), - config.getCloudWatchClientConfiguration()); + cloudWatchClient = (AmazonCloudWatch) createClient(AmazonCloudWatchClientBuilder.standard(), + 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()); - cloudWatchClient.setRegion(region); - LOG.debug("The region of Amazon CloudWatch client has been set to " + 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()); + setRegionForClient((AmazonWebServiceClient) cloudWatchClient, config.getRegionName()); + setRegionForClient((AmazonWebServiceClient) kinesisClient, config.getRegionName()); + setRegionForClient((AmazonWebServiceClient) dynamoDBClient, config.getRegionName()); } // If a dynamoDB endpoint was explicitly specified, use it to set the DynamoDB endpoint. if (config.getDynamoDBEndpoint() != null) { - dynamoDBClient.setEndpoint(config.getDynamoDBEndpoint()); - LOG.debug("The endpoint of Amazon DynamoDB client has been set to " + config.getDynamoDBEndpoint()); + setEndpointForClient((AmazonWebServiceClient) dynamoDBClient, 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()); - 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()); - } + setEndpointForClient((AmazonWebServiceClient) kinesisClient, 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(); + } + } } diff --git a/src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/KinesisClientLibConfigurationTest.java b/src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/KinesisClientLibConfigurationTest.java index cccbcb30..d30c9e57 100644 --- a/src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/KinesisClientLibConfigurationTest.java +++ b/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(4)).getKinesisEndpoint(); + Mockito.verify(kclConfig, Mockito.times(5)).getRegionName(); + 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(3)).getKinesisEndpoint(); + Mockito.verify(kclConfig, Mockito.times(2)).getRegionName(); + Mockito.verify(kclConfig, Mockito.times(2)).getKinesisEndpoint(); } @Test diff --git a/src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/WorkerTest.java b/src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/WorkerTest.java index 037a54b2..06dffa50 100644 --- a/src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/WorkerTest.java +++ b/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,