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 cca09c44..3c39aa33 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 @@ -33,6 +33,8 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import com.amazonaws.services.kinesis.clientlibrary.proxies.IKinesisProxy; +import com.amazonaws.services.kinesis.leases.interfaces.ILeaseManager; + import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; @@ -1076,6 +1078,7 @@ public class Worker implements Runnable { private AmazonDynamoDB dynamoDBClient; private AmazonCloudWatch cloudWatchClient; private IMetricsFactory metricsFactory; + private ILeaseManager leaseManager; private ExecutorService execService; private ShardPrioritization shardPrioritization; private IKinesisProxy kinesisProxy; @@ -1173,6 +1176,18 @@ public class Worker implements Runnable { return this; } + /** + * Set the lease manager. + * + * @param leaseManager + * Lease manager used to manage shard leases + * @return A reference to this updated object so that method calls can be chained together. + */ + public Builder leaseManager(ILeaseManager leaseManager) { + this.leaseManager = leaseManager; + return this; + } + /** * Set the executor service for processing records. * @@ -1273,6 +1288,9 @@ public class Worker implements Runnable { if (metricsFactory == null) { metricsFactory = getMetricsFactory(cloudWatchClient, config); } + if (leaseManager == null) { + leaseManager = new KinesisClientLeaseManager(config.getTableName(), dynamoDBClient); + } if (shardPrioritization == null) { shardPrioritization = new ParentsFirstShardPrioritization(1); } @@ -1294,8 +1312,7 @@ public class Worker implements Runnable { config.getShardSyncIntervalMillis(), config.shouldCleanupLeasesUponShardCompletion(), null, - new KinesisClientLibLeaseCoordinator(new KinesisClientLeaseManager(config.getTableName(), - dynamoDBClient), + new KinesisClientLibLeaseCoordinator(leaseManager, config.getWorkerIdentifier(), config.getFailoverTimeMillis(), config.getEpsilonMillis(),