Flattening LeaseManagement to DynamoDBLeaseManagement.

This commit is contained in:
Sahil Palvia 2018-04-23 14:36:15 -07:00
parent c256f75cc8
commit 12a173231d
60 changed files with 2225 additions and 2719 deletions

View file

@ -15,15 +15,13 @@
package software.amazon.kinesis.checkpoint; package software.amazon.kinesis.checkpoint;
import software.amazon.kinesis.leases.LeaseManager;
import software.amazon.kinesis.leases.KinesisClientLease;
import software.amazon.kinesis.leases.LeaseCoordinator; import software.amazon.kinesis.leases.LeaseCoordinator;
import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.processor.Checkpointer; import software.amazon.kinesis.processor.Checkpointer;
/** /**
* *
*/ */
public interface CheckpointFactory { public interface CheckpointFactory {
Checkpointer createCheckpointer(LeaseCoordinator<KinesisClientLease> leaseCoordinator, Checkpointer createCheckpointer(LeaseCoordinator leaseCoordinator, LeaseRefresher leaseRefresher);
LeaseManager<KinesisClientLease> leaseManager);
} }

View file

@ -17,9 +17,8 @@ package software.amazon.kinesis.checkpoint;
import lombok.Data; import lombok.Data;
import lombok.NonNull; import lombok.NonNull;
import software.amazon.kinesis.leases.LeaseManager;
import software.amazon.kinesis.leases.KinesisClientLease;
import software.amazon.kinesis.leases.LeaseCoordinator; import software.amazon.kinesis.leases.LeaseCoordinator;
import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.metrics.IMetricsFactory; import software.amazon.kinesis.metrics.IMetricsFactory;
import software.amazon.kinesis.processor.Checkpointer; import software.amazon.kinesis.processor.Checkpointer;
@ -32,9 +31,9 @@ public class DynamoDBCheckpointFactory implements CheckpointFactory {
private final IMetricsFactory metricsFactory; private final IMetricsFactory metricsFactory;
@Override @Override
public Checkpointer createCheckpointer(final LeaseCoordinator<KinesisClientLease> leaseLeaseCoordinator, public Checkpointer createCheckpointer(final LeaseCoordinator leaseLeaseCoordinator,
final LeaseManager<KinesisClientLease> leaseManager) { final LeaseRefresher leaseRefresher) {
return new DynamoDBCheckpointer(leaseLeaseCoordinator, leaseManager, metricsFactory); return new DynamoDBCheckpointer(leaseLeaseCoordinator, leaseRefresher, metricsFactory);
} }
} }

View file

@ -28,9 +28,9 @@ import com.google.common.annotations.VisibleForTesting;
import lombok.NonNull; import lombok.NonNull;
import lombok.RequiredArgsConstructor; import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.leases.LeaseManager; import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.KinesisClientLease;
import software.amazon.kinesis.leases.LeaseCoordinator; import software.amazon.kinesis.leases.LeaseCoordinator;
import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException; import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
@ -45,9 +45,9 @@ import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
@Slf4j @Slf4j
public class DynamoDBCheckpointer implements Checkpointer { public class DynamoDBCheckpointer implements Checkpointer {
@NonNull @NonNull
private final LeaseCoordinator<KinesisClientLease> leaseCoordinator; private final LeaseCoordinator leaseCoordinator;
@NonNull @NonNull
private final LeaseManager<KinesisClientLease> leaseManager; private final LeaseRefresher leaseRefresher;
@NonNull @NonNull
private final IMetricsFactory metricsFactory; private final IMetricsFactory metricsFactory;
@ -73,7 +73,7 @@ public class DynamoDBCheckpointer implements Checkpointer {
@Override @Override
public ExtendedSequenceNumber getCheckpoint(final String shardId) throws KinesisClientLibException { public ExtendedSequenceNumber getCheckpoint(final String shardId) throws KinesisClientLibException {
try { try {
return leaseManager.getLease(shardId).getCheckpoint(); return leaseRefresher.getLease(shardId).checkpoint();
} catch (DependencyException | InvalidStateException | ProvisionedThroughputException e) { } catch (DependencyException | InvalidStateException | ProvisionedThroughputException e) {
String message = "Unable to fetch checkpoint for shardId " + shardId; String message = "Unable to fetch checkpoint for shardId " + shardId;
log.error(message, e); log.error(message, e);
@ -84,8 +84,8 @@ public class DynamoDBCheckpointer implements Checkpointer {
@Override @Override
public Checkpoint getCheckpointObject(final String shardId) throws KinesisClientLibException { public Checkpoint getCheckpointObject(final String shardId) throws KinesisClientLibException {
try { try {
KinesisClientLease lease = leaseManager.getLease(shardId); Lease lease = leaseRefresher.getLease(shardId);
return new Checkpoint(lease.getCheckpoint(), lease.getPendingCheckpoint()); return new Checkpoint(lease.checkpoint(), lease.pendingCheckpoint());
} catch (DependencyException | InvalidStateException | ProvisionedThroughputException e) { } catch (DependencyException | InvalidStateException | ProvisionedThroughputException e) {
String message = "Unable to fetch checkpoint for shardId " + shardId; String message = "Unable to fetch checkpoint for shardId " + shardId;
log.error(message, e); log.error(message, e);
@ -117,30 +117,30 @@ public class DynamoDBCheckpointer implements Checkpointer {
@VisibleForTesting @VisibleForTesting
public boolean setCheckpoint(String shardId, ExtendedSequenceNumber checkpoint, UUID concurrencyToken) public boolean setCheckpoint(String shardId, ExtendedSequenceNumber checkpoint, UUID concurrencyToken)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
KinesisClientLease lease = leaseCoordinator.getCurrentlyHeldLease(shardId); Lease lease = leaseCoordinator.getCurrentlyHeldLease(shardId);
if (lease == null) { if (lease == null) {
log.info("Worker {} could not update checkpoint for shard {} because it does not hold the lease", log.info("Worker {} could not update checkpoint for shard {} because it does not hold the lease",
leaseCoordinator.getWorkerIdentifier(), shardId); leaseCoordinator.workerIdentifier(), shardId);
return false; return false;
} }
lease.setCheckpoint(checkpoint); lease.checkpoint(checkpoint);
lease.setPendingCheckpoint(null); lease.pendingCheckpoint(null);
lease.setOwnerSwitchesSinceCheckpoint(0L); lease.ownerSwitchesSinceCheckpoint(0L);
return leaseCoordinator.updateLease(lease, concurrencyToken); return leaseCoordinator.updateLease(lease, concurrencyToken);
} }
boolean prepareCheckpoint(String shardId, ExtendedSequenceNumber pendingCheckpoint, UUID concurrencyToken) boolean prepareCheckpoint(String shardId, ExtendedSequenceNumber pendingCheckpoint, UUID concurrencyToken)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
KinesisClientLease lease = leaseCoordinator.getCurrentlyHeldLease(shardId); Lease lease = leaseCoordinator.getCurrentlyHeldLease(shardId);
if (lease == null) { if (lease == null) {
log.info("Worker {} could not prepare checkpoint for shard {} because it does not hold the lease", log.info("Worker {} could not prepare checkpoint for shard {} because it does not hold the lease",
leaseCoordinator.getWorkerIdentifier(), shardId); leaseCoordinator.workerIdentifier(), shardId);
return false; return false;
} }
lease.setPendingCheckpoint(Objects.requireNonNull(pendingCheckpoint, "pendingCheckpoint should not be null")); lease.pendingCheckpoint(Objects.requireNonNull(pendingCheckpoint, "pendingCheckpoint should not be null"));
return leaseCoordinator.updateLease(lease, concurrencyToken); return leaseCoordinator.updateLease(lease, concurrencyToken);
} }
} }

View file

@ -38,11 +38,12 @@ import lombok.NonNull;
import lombok.experimental.Accessors; import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.checkpoint.CheckpointConfig; import software.amazon.kinesis.checkpoint.CheckpointConfig;
import software.amazon.kinesis.leases.LeaseManager; import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseCoordinator;
import software.amazon.kinesis.leases.KinesisClientLease; import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.KinesisClientLibLeaseCoordinator; import software.amazon.kinesis.leases.LeaseCoordinator;
import software.amazon.kinesis.leases.LeaseManagementConfig; import software.amazon.kinesis.leases.LeaseManagementConfig;
import software.amazon.kinesis.leases.LeaseManagerProxy; import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.ShardDetector;
import software.amazon.kinesis.leases.ShardInfo; import software.amazon.kinesis.leases.ShardInfo;
import software.amazon.kinesis.leases.ShardPrioritization; import software.amazon.kinesis.leases.ShardPrioritization;
import software.amazon.kinesis.leases.ShardSyncTask; import software.amazon.kinesis.leases.ShardSyncTask;
@ -60,9 +61,9 @@ import software.amazon.kinesis.metrics.MetricsCollectingTaskDecorator;
import software.amazon.kinesis.metrics.MetricsConfig; import software.amazon.kinesis.metrics.MetricsConfig;
import software.amazon.kinesis.metrics.MetricsLevel; import software.amazon.kinesis.metrics.MetricsLevel;
import software.amazon.kinesis.processor.Checkpointer; import software.amazon.kinesis.processor.Checkpointer;
import software.amazon.kinesis.processor.ShutdownNotificationAware;
import software.amazon.kinesis.processor.ProcessorConfig; import software.amazon.kinesis.processor.ProcessorConfig;
import software.amazon.kinesis.processor.ProcessorFactory; import software.amazon.kinesis.processor.ProcessorFactory;
import software.amazon.kinesis.processor.ShutdownNotificationAware;
import software.amazon.kinesis.retrieval.RetrievalConfig; import software.amazon.kinesis.retrieval.RetrievalConfig;
/** /**
@ -91,7 +92,7 @@ public class Scheduler implements Runnable {
private final long parentShardPollIntervalMillis; private final long parentShardPollIntervalMillis;
private final ExecutorService executorService; private final ExecutorService executorService;
// private final GetRecordsRetrievalStrategy getRecordsRetrievalStrategy; // private final GetRecordsRetrievalStrategy getRecordsRetrievalStrategy;
private final KinesisClientLibLeaseCoordinator leaseCoordinator; private final LeaseCoordinator leaseCoordinator;
private final ShardSyncTaskManager shardSyncTaskManager; private final ShardSyncTaskManager shardSyncTaskManager;
private final ShardPrioritization shardPrioritization; private final ShardPrioritization shardPrioritization;
private final boolean cleanupLeasesUponShardCompletion; private final boolean cleanupLeasesUponShardCompletion;
@ -108,8 +109,8 @@ public class Scheduler implements Runnable {
private final String streamName; private final String streamName;
private final long listShardsBackoffTimeMillis; private final long listShardsBackoffTimeMillis;
private final int maxListShardsRetryAttempts; private final int maxListShardsRetryAttempts;
private final LeaseManager<KinesisClientLease> leaseManager; private final LeaseRefresher leaseRefresher;
private final LeaseManagerProxy leaseManagerProxy; private final ShardDetector shardDetector;
private final boolean ignoreUnexpetedChildShards; private final boolean ignoreUnexpetedChildShards;
// Holds consumers for shards the worker is currently tracking. Key is shard // Holds consumers for shards the worker is currently tracking. Key is shard
@ -143,15 +144,14 @@ public class Scheduler implements Runnable {
this.retrievalConfig = retrievalConfig; this.retrievalConfig = retrievalConfig;
this.applicationName = this.coordinatorConfig.applicationName(); this.applicationName = this.coordinatorConfig.applicationName();
this.leaseCoordinator = this.leaseCoordinator = this.leaseManagementConfig.leaseManagementFactory().createLeaseCoordinator();
this.leaseManagementConfig.leaseManagementFactory().createKinesisClientLibLeaseCoordinator(); this.leaseRefresher = this.leaseCoordinator.leaseRefresher();
this.leaseManager = this.leaseCoordinator.leaseManager();
// //
// TODO: Figure out what to do with lease manage <=> checkpoint relationship // TODO: Figure out what to do with lease manage <=> checkpoint relationship
// //
this.checkpoint = this.checkpointConfig.checkpointFactory().createCheckpointer(this.leaseCoordinator, this.checkpoint = this.checkpointConfig.checkpointFactory().createCheckpointer(this.leaseCoordinator,
this.leaseManager); this.leaseRefresher);
this.idleTimeInMilliseconds = this.retrievalConfig.idleTimeBetweenReadsInMillis(); this.idleTimeInMilliseconds = this.retrievalConfig.idleTimeBetweenReadsInMillis();
this.parentShardPollIntervalMillis = this.coordinatorConfig.parentShardPollIntervalMillis(); this.parentShardPollIntervalMillis = this.coordinatorConfig.parentShardPollIntervalMillis();
@ -175,7 +175,7 @@ public class Scheduler implements Runnable {
this.streamName = this.retrievalConfig.streamName(); this.streamName = this.retrievalConfig.streamName();
this.listShardsBackoffTimeMillis = this.retrievalConfig.listShardsBackoffTimeInMillis(); this.listShardsBackoffTimeMillis = this.retrievalConfig.listShardsBackoffTimeInMillis();
this.maxListShardsRetryAttempts = this.retrievalConfig.maxListShardsRetryAttempts(); this.maxListShardsRetryAttempts = this.retrievalConfig.maxListShardsRetryAttempts();
this.leaseManagerProxy = this.shardSyncTaskManager.leaseManagerProxy(); this.shardDetector = this.shardSyncTaskManager.shardDetector();
this.ignoreUnexpetedChildShards = this.leaseManagementConfig.ignoreUnexpectedChildShards(); this.ignoreUnexpetedChildShards = this.leaseManagementConfig.ignoreUnexpectedChildShards();
} }
@ -217,9 +217,9 @@ public class Scheduler implements Runnable {
TaskResult result = null; TaskResult result = null;
if (!skipShardSyncAtWorkerInitializationIfLeasesExist if (!skipShardSyncAtWorkerInitializationIfLeasesExist
|| leaseManager.isLeaseTableEmpty()) { || leaseRefresher.isLeaseTableEmpty()) {
log.info("Syncing Kinesis shard info"); log.info("Syncing Kinesis shard info");
ShardSyncTask shardSyncTask = new ShardSyncTask(leaseManagerProxy, leaseManager, initialPosition, ShardSyncTask shardSyncTask = new ShardSyncTask(shardDetector, leaseRefresher, initialPosition,
cleanupLeasesUponShardCompletion, ignoreUnexpetedChildShards, 0L); cleanupLeasesUponShardCompletion, ignoreUnexpetedChildShards, 0L);
result = new MetricsCollectingTaskDecorator(shardSyncTask, metricsFactory).call(); result = new MetricsCollectingTaskDecorator(shardSyncTask, metricsFactory).call();
} else { } else {
@ -398,7 +398,7 @@ public class Scheduler implements Runnable {
// //
leaseCoordinator.stopLeaseTaker(); leaseCoordinator.stopLeaseTaker();
Collection<KinesisClientLease> leases = leaseCoordinator.getAssignments(); Collection<Lease> leases = leaseCoordinator.getAssignments();
if (leases == null || leases.isEmpty()) { if (leases == null || leases.isEmpty()) {
// //
// If there are no leases notification is already completed, but we still need to shutdown the worker. // If there are no leases notification is already completed, but we still need to shutdown the worker.
@ -408,10 +408,10 @@ public class Scheduler implements Runnable {
} }
CountDownLatch shutdownCompleteLatch = new CountDownLatch(leases.size()); CountDownLatch shutdownCompleteLatch = new CountDownLatch(leases.size());
CountDownLatch notificationCompleteLatch = new CountDownLatch(leases.size()); CountDownLatch notificationCompleteLatch = new CountDownLatch(leases.size());
for (KinesisClientLease lease : leases) { for (Lease lease : leases) {
ShutdownNotification shutdownNotification = new ShardConsumerShutdownNotification(leaseCoordinator, ShutdownNotification shutdownNotification = new ShardConsumerShutdownNotification(leaseCoordinator,
lease, notificationCompleteLatch, shutdownCompleteLatch); lease, notificationCompleteLatch, shutdownCompleteLatch);
ShardInfo shardInfo = KinesisClientLibLeaseCoordinator.convertLeaseToAssignment(lease); ShardInfo shardInfo = DynamoDBLeaseCoordinator.convertLeaseToAssignment(lease);
ShardConsumer consumer = shardInfoShardConsumerMap.get(shardInfo); ShardConsumer consumer = shardInfoShardConsumerMap.get(shardInfo);
if (consumer != null) { if (consumer != null) {
consumer.notifyShutdownRequested(shutdownNotification); consumer.notifyShutdownRequested(shutdownNotification);
@ -531,7 +531,7 @@ public class Scheduler implements Runnable {
@NonNull final ProcessorFactory processorFactory) { @NonNull final ProcessorFactory processorFactory) {
return new ShardConsumer(shardInfo, return new ShardConsumer(shardInfo,
streamName, streamName,
leaseManager, leaseRefresher,
executorService, executorService,
retrievalConfig.retrievalFactory().createGetRecordsCache(shardInfo, metricsFactory), retrievalConfig.retrievalFactory().createGetRecordsCache(shardInfo, metricsFactory),
processorFactory.createRecordProcessor(), processorFactory.createRecordProcessor(),
@ -551,7 +551,7 @@ public class Scheduler implements Runnable {
initialPosition, initialPosition,
cleanupLeasesUponShardCompletion, cleanupLeasesUponShardCompletion,
ignoreUnexpetedChildShards, ignoreUnexpetedChildShards,
leaseManagerProxy, shardDetector,
metricsFactory); metricsFactory);
} }

View file

@ -1,194 +0,0 @@
/*
* Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Amazon Software License (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/asl/
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*/
package software.amazon.kinesis.leases;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import com.amazonaws.services.dynamodbv2.model.AttributeAction;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.AttributeValue;
import com.amazonaws.services.dynamodbv2.model.AttributeValueUpdate;
import com.amazonaws.services.dynamodbv2.model.ExpectedAttributeValue;
import com.amazonaws.services.dynamodbv2.model.KeySchemaElement;
import com.amazonaws.services.dynamodbv2.model.KeyType;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
/**
* An implementation of ILeaseSerializer for basic Lease objects. Can also instantiate subclasses of Lease so that
* LeaseSerializer can be decorated by other classes if you need to add fields to leases.
*/
public class DynamoDBLeaseSerializer implements LeaseSerializer<Lease> {
public final String LEASE_KEY_KEY = "leaseKey";
public final String LEASE_OWNER_KEY = "leaseOwner";
public final String LEASE_COUNTER_KEY = "leaseCounter";
public final Class<? extends Lease> clazz;
public DynamoDBLeaseSerializer() {
this.clazz = Lease.class;
}
public DynamoDBLeaseSerializer(Class<? extends Lease> clazz) {
this.clazz = clazz;
}
@Override
public Map<String, AttributeValue> toDynamoRecord(Lease lease) {
Map<String, AttributeValue> result = new HashMap<String, AttributeValue>();
result.put(LEASE_KEY_KEY, DynamoUtils.createAttributeValue(lease.getLeaseKey()));
result.put(LEASE_COUNTER_KEY, DynamoUtils.createAttributeValue(lease.getLeaseCounter()));
if (lease.getLeaseOwner() != null) {
result.put(LEASE_OWNER_KEY, DynamoUtils.createAttributeValue(lease.getLeaseOwner()));
}
return result;
}
@Override
public Lease fromDynamoRecord(Map<String, AttributeValue> dynamoRecord) {
Lease result;
try {
result = clazz.newInstance();
} catch (InstantiationException e) {
throw new RuntimeException(e);
} catch (IllegalAccessException e) {
throw new RuntimeException(e);
}
result.setLeaseKey(DynamoUtils.safeGetString(dynamoRecord, LEASE_KEY_KEY));
result.setLeaseOwner(DynamoUtils.safeGetString(dynamoRecord, LEASE_OWNER_KEY));
result.setLeaseCounter(DynamoUtils.safeGetLong(dynamoRecord, LEASE_COUNTER_KEY));
return result;
}
@Override
public Map<String, AttributeValue> getDynamoHashKey(String leaseKey) {
Map<String, AttributeValue> result = new HashMap<String, AttributeValue>();
result.put(LEASE_KEY_KEY, DynamoUtils.createAttributeValue(leaseKey));
return result;
}
@Override
public Map<String, AttributeValue> getDynamoHashKey(Lease lease) {
return getDynamoHashKey(lease.getLeaseKey());
}
@Override
public Map<String, ExpectedAttributeValue> getDynamoLeaseCounterExpectation(Lease lease) {
return getDynamoLeaseCounterExpectation(lease.getLeaseCounter());
}
public Map<String, ExpectedAttributeValue> getDynamoLeaseCounterExpectation(Long leaseCounter) {
Map<String, ExpectedAttributeValue> result = new HashMap<String, ExpectedAttributeValue>();
ExpectedAttributeValue eav = new ExpectedAttributeValue(DynamoUtils.createAttributeValue(leaseCounter));
result.put(LEASE_COUNTER_KEY, eav);
return result;
}
@Override
public Map<String, ExpectedAttributeValue> getDynamoLeaseOwnerExpectation(Lease lease) {
Map<String, ExpectedAttributeValue> result = new HashMap<String, ExpectedAttributeValue>();
ExpectedAttributeValue eav = null;
if (lease.getLeaseOwner() == null) {
eav = new ExpectedAttributeValue(false);
} else {
eav = new ExpectedAttributeValue(DynamoUtils.createAttributeValue(lease.getLeaseOwner()));
}
result.put(LEASE_OWNER_KEY, eav);
return result;
}
@Override
public Map<String, ExpectedAttributeValue> getDynamoNonexistantExpectation() {
Map<String, ExpectedAttributeValue> result = new HashMap<String, ExpectedAttributeValue>();
ExpectedAttributeValue expectedAV = new ExpectedAttributeValue(false);
result.put(LEASE_KEY_KEY, expectedAV);
return result;
}
@Override
public Map<String, AttributeValueUpdate> getDynamoLeaseCounterUpdate(Lease lease) {
return getDynamoLeaseCounterUpdate(lease.getLeaseCounter());
}
public Map<String, AttributeValueUpdate> getDynamoLeaseCounterUpdate(Long leaseCounter) {
Map<String, AttributeValueUpdate> result = new HashMap<String, AttributeValueUpdate>();
AttributeValueUpdate avu =
new AttributeValueUpdate(DynamoUtils.createAttributeValue(leaseCounter + 1), AttributeAction.PUT);
result.put(LEASE_COUNTER_KEY, avu);
return result;
}
@Override
public Map<String, AttributeValueUpdate> getDynamoTakeLeaseUpdate(Lease lease, String owner) {
Map<String, AttributeValueUpdate> result = new HashMap<String, AttributeValueUpdate>();
result.put(LEASE_OWNER_KEY, new AttributeValueUpdate(DynamoUtils.createAttributeValue(owner),
AttributeAction.PUT));
return result;
}
@Override
public Map<String, AttributeValueUpdate> getDynamoEvictLeaseUpdate(Lease lease) {
Map<String, AttributeValueUpdate> result = new HashMap<String, AttributeValueUpdate>();
result.put(LEASE_OWNER_KEY, new AttributeValueUpdate(null, AttributeAction.DELETE));
return result;
}
@Override
public Map<String, AttributeValueUpdate> getDynamoUpdateLeaseUpdate(Lease lease) {
// There is no application-specific data in Lease - just return a map that increments the counter.
return new HashMap<String, AttributeValueUpdate>();
}
@Override
public Collection<KeySchemaElement> getKeySchema() {
List<KeySchemaElement> keySchema = new ArrayList<KeySchemaElement>();
keySchema.add(new KeySchemaElement().withAttributeName(LEASE_KEY_KEY).withKeyType(KeyType.HASH));
return keySchema;
}
@Override
public Collection<AttributeDefinition> getAttributeDefinitions() {
List<AttributeDefinition> definitions = new ArrayList<AttributeDefinition>();
definitions.add(new AttributeDefinition().withAttributeName(LEASE_KEY_KEY)
.withAttributeType(ScalarAttributeType.S));
return definitions;
}
}

View file

@ -1,79 +0,0 @@
/*
* Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Amazon Software License (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/asl/
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*/
package software.amazon.kinesis.leases;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
/**
* An implementation of LeaseManager for the KinesisClientLibrary - takeLease updates the ownerSwitchesSinceCheckpoint field.
*/
public class KinesisClientDynamoDBLeaseManager extends DynamoDBLeaseManager<KinesisClientLease> implements KinesisClientLeaseManager {
/**
* Constructor.
*
* @param table Leases table
* @param dynamoDBClient DynamoDB client to use
*/
public KinesisClientDynamoDBLeaseManager(String table, AmazonDynamoDB dynamoDBClient) {
this(table, dynamoDBClient, false);
}
/**
* Constructor for integration tests - see comment on superclass for documentation on setting the consistentReads
* flag.
*
* @param table leases table
* @param dynamoDBClient DynamoDB client to use
* @param consistentReads true if we want consistent reads for testing purposes.
*/
public KinesisClientDynamoDBLeaseManager(String table, AmazonDynamoDB dynamoDBClient, boolean consistentReads) {
super(table, dynamoDBClient, new KinesisClientLeaseSerializer(), consistentReads);
}
/**
* {@inheritDoc}
*/
@Override
public boolean takeLease(KinesisClientLease lease, String newOwner)
throws DependencyException, InvalidStateException, ProvisionedThroughputException {
String oldOwner = lease.getLeaseOwner();
boolean result = super.takeLease(lease, newOwner);
if (oldOwner != null && !oldOwner.equals(newOwner)) {
lease.setOwnerSwitchesSinceCheckpoint(lease.getOwnerSwitchesSinceCheckpoint() + 1);
}
return result;
}
/**
* {@inheritDoc}
*/
@Override
public ExtendedSequenceNumber getCheckpoint(String shardId)
throws ProvisionedThroughputException, InvalidStateException, DependencyException {
ExtendedSequenceNumber checkpoint = null;
KinesisClientLease lease = getLease(shardId);
if (lease != null) {
checkpoint = lease.getCheckpoint();
}
return checkpoint;
}
}

View file

@ -1,207 +0,0 @@
/*
* Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Amazon Software License (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/asl/
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*/
package software.amazon.kinesis.leases;
import java.util.Collection;
import java.util.HashSet;
import java.util.Set;
import java.util.UUID;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
/**
* A Lease subclass containing KinesisClientLibrary related fields for checkpoints.
*/
public class KinesisClientLease extends Lease {
private ExtendedSequenceNumber checkpoint;
private ExtendedSequenceNumber pendingCheckpoint;
private Long ownerSwitchesSinceCheckpoint = 0L;
private Set<String> parentShardIds = new HashSet<String>();
public KinesisClientLease() {
}
public KinesisClientLease(KinesisClientLease other) {
super(other);
this.checkpoint = other.getCheckpoint();
this.pendingCheckpoint = other.getPendingCheckpoint();
this.ownerSwitchesSinceCheckpoint = other.getOwnerSwitchesSinceCheckpoint();
this.parentShardIds.addAll(other.getParentShardIds());
}
KinesisClientLease(String leaseKey, String leaseOwner, Long leaseCounter, UUID concurrencyToken,
Long lastCounterIncrementNanos, ExtendedSequenceNumber checkpoint, ExtendedSequenceNumber pendingCheckpoint,
Long ownerSwitchesSinceCheckpoint, Set<String> parentShardIds) {
super(leaseKey, leaseOwner, leaseCounter, concurrencyToken, lastCounterIncrementNanos);
this.checkpoint = checkpoint;
this.pendingCheckpoint = pendingCheckpoint;
this.ownerSwitchesSinceCheckpoint = ownerSwitchesSinceCheckpoint;
this.parentShardIds.addAll(parentShardIds);
}
/**
* {@inheritDoc}
*/
@Override
public <T extends Lease> void update(T other) {
super.update(other);
if (!(other instanceof KinesisClientLease)) {
throw new IllegalArgumentException("Must pass KinesisClientLease object to KinesisClientLease.update(Lease)");
}
KinesisClientLease casted = (KinesisClientLease) other;
setOwnerSwitchesSinceCheckpoint(casted.ownerSwitchesSinceCheckpoint);
setCheckpoint(casted.checkpoint);
setPendingCheckpoint(casted.pendingCheckpoint);
setParentShardIds(casted.parentShardIds);
}
/**
* @return most recently application-supplied checkpoint value. During fail over, the new worker will pick up after
* the old worker's last checkpoint.
*/
public ExtendedSequenceNumber getCheckpoint() {
return checkpoint;
}
/**
* @return pending checkpoint, possibly null.
*/
public ExtendedSequenceNumber getPendingCheckpoint() {
return pendingCheckpoint;
}
/**
* @return count of distinct lease holders between checkpoints.
*/
public Long getOwnerSwitchesSinceCheckpoint() {
return ownerSwitchesSinceCheckpoint;
}
/**
* @return shardIds that parent this lease. Used for resharding.
*/
public Set<String> getParentShardIds() {
return new HashSet<String>(parentShardIds);
}
/**
* Sets checkpoint.
*
* @param checkpoint may not be null
*/
public void setCheckpoint(ExtendedSequenceNumber checkpoint) {
verifyNotNull(checkpoint, "Checkpoint should not be null");
this.checkpoint = checkpoint;
}
/**
* Sets pending checkpoint.
*
* @param pendingCheckpoint can be null
*/
public void setPendingCheckpoint(ExtendedSequenceNumber pendingCheckpoint) {
this.pendingCheckpoint = pendingCheckpoint;
}
/**
* Sets ownerSwitchesSinceCheckpoint.
*
* @param ownerSwitchesSinceCheckpoint may not be null
*/
public void setOwnerSwitchesSinceCheckpoint(Long ownerSwitchesSinceCheckpoint) {
verifyNotNull(ownerSwitchesSinceCheckpoint, "ownerSwitchesSinceCheckpoint should not be null");
this.ownerSwitchesSinceCheckpoint = ownerSwitchesSinceCheckpoint;
}
/**
* Sets parentShardIds.
*
* @param parentShardIds may not be null
*/
public void setParentShardIds(Collection<String> parentShardIds) {
verifyNotNull(parentShardIds, "parentShardIds should not be null");
this.parentShardIds.clear();
this.parentShardIds.addAll(parentShardIds);
}
private void verifyNotNull(Object object, String message) {
if (object == null) {
throw new IllegalArgumentException(message);
}
}
@Override
public int hashCode() {
final int prime = 31;
int result = super.hashCode();
result = prime * result + ((checkpoint == null) ? 0 : checkpoint.hashCode());
result = pendingCheckpoint == null ? result : prime * result + pendingCheckpoint.hashCode();
result =
prime * result + ((ownerSwitchesSinceCheckpoint == null) ? 0 : ownerSwitchesSinceCheckpoint.hashCode());
result = prime * result + ((parentShardIds == null) ? 0 : parentShardIds.hashCode());
return result;
}
@Override
public boolean equals(Object obj) {
if (this == obj)
return true;
if (!super.equals(obj))
return false;
if (getClass() != obj.getClass())
return false;
KinesisClientLease other = (KinesisClientLease) obj;
if (checkpoint == null) {
if (other.checkpoint != null)
return false;
} else if (!checkpoint.equals(other.checkpoint))
return false;
if (pendingCheckpoint == null) {
if (other.pendingCheckpoint != null)
return false;
} else if (!pendingCheckpoint.equals(other.pendingCheckpoint))
return false;
if (ownerSwitchesSinceCheckpoint == null) {
if (other.ownerSwitchesSinceCheckpoint != null)
return false;
} else if (!ownerSwitchesSinceCheckpoint.equals(other.ownerSwitchesSinceCheckpoint))
return false;
if (parentShardIds == null) {
if (other.parentShardIds != null)
return false;
} else if (!parentShardIds.equals(other.parentShardIds))
return false;
return true;
}
/**
* Returns a deep copy of this object. Type-unsafe - there aren't good mechanisms for copy-constructing generics.
*
* @return A deep copy of this object.
*/
@Override
@SuppressWarnings("unchecked")
public <T extends Lease> T copy() {
return (T) new KinesisClientLease(this);
}
}

View file

@ -1,40 +0,0 @@
/*
* Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Amazon Software License (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/asl/
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*/
package software.amazon.kinesis.leases;
import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
/**
* A decoration of ILeaseManager that adds methods to get/update checkpoints.
*/
public interface KinesisClientLeaseManager extends LeaseManager<KinesisClientLease> {
/**
* Gets the current checkpoint of the shard. This is useful in the resharding use case
* where we will wait for the parent shard to complete before starting on the records from a child shard.
*
* @param shardId Checkpoint of this shard will be returned
* @return Checkpoint of this shard, or null if the shard record doesn't exist.
*
* @throws ProvisionedThroughputException if DynamoDB update fails due to lack of capacity
* @throws InvalidStateException if lease table does not exist
* @throws DependencyException if DynamoDB update fails in an unexpected way
*/
ExtendedSequenceNumber getCheckpoint(String shardId) throws ProvisionedThroughputException, InvalidStateException, DependencyException;
}

View file

@ -1,166 +0,0 @@
/*
* Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Amazon Software License (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/asl/
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*/
package software.amazon.kinesis.leases;
import java.util.Collection;
import java.util.Map;
import com.amazonaws.services.dynamodbv2.model.AttributeAction;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.AttributeValue;
import com.amazonaws.services.dynamodbv2.model.AttributeValueUpdate;
import com.amazonaws.services.dynamodbv2.model.ExpectedAttributeValue;
import com.amazonaws.services.dynamodbv2.model.KeySchemaElement;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
import com.google.common.base.Strings;
/**
* An implementation of ILeaseSerializer for KinesisClientLease objects.
*/
public class KinesisClientLeaseSerializer implements LeaseSerializer<KinesisClientLease> {
private static final String OWNER_SWITCHES_KEY = "ownerSwitchesSinceCheckpoint";
private static final String CHECKPOINT_SEQUENCE_NUMBER_KEY = "checkpoint";
private static final String CHECKPOINT_SUBSEQUENCE_NUMBER_KEY = "checkpointSubSequenceNumber";
private static final String PENDING_CHECKPOINT_SEQUENCE_KEY = "pendingCheckpoint";
private static final String PENDING_CHECKPOINT_SUBSEQUENCE_KEY = "pendingCheckpointSubSequenceNumber";
public final String PARENT_SHARD_ID_KEY = "parentShardId";
private final DynamoDBLeaseSerializer baseSerializer = new DynamoDBLeaseSerializer(KinesisClientLease.class);
@Override
public Map<String, AttributeValue> toDynamoRecord(KinesisClientLease lease) {
Map<String, AttributeValue> result = baseSerializer.toDynamoRecord(lease);
result.put(OWNER_SWITCHES_KEY, DynamoUtils.createAttributeValue(lease.getOwnerSwitchesSinceCheckpoint()));
result.put(CHECKPOINT_SEQUENCE_NUMBER_KEY, DynamoUtils.createAttributeValue(lease.getCheckpoint().getSequenceNumber()));
result.put(CHECKPOINT_SUBSEQUENCE_NUMBER_KEY, DynamoUtils.createAttributeValue(lease.getCheckpoint().getSubSequenceNumber()));
if (lease.getParentShardIds() != null && !lease.getParentShardIds().isEmpty()) {
result.put(PARENT_SHARD_ID_KEY, DynamoUtils.createAttributeValue(lease.getParentShardIds()));
}
if (lease.getPendingCheckpoint() != null && !lease.getPendingCheckpoint().getSequenceNumber().isEmpty()) {
result.put(PENDING_CHECKPOINT_SEQUENCE_KEY, DynamoUtils.createAttributeValue(lease.getPendingCheckpoint().getSequenceNumber()));
result.put(PENDING_CHECKPOINT_SUBSEQUENCE_KEY, DynamoUtils.createAttributeValue(lease.getPendingCheckpoint().getSubSequenceNumber()));
}
return result;
}
@Override
public KinesisClientLease fromDynamoRecord(Map<String, AttributeValue> dynamoRecord) {
KinesisClientLease result = (KinesisClientLease) baseSerializer.fromDynamoRecord(dynamoRecord);
result.setOwnerSwitchesSinceCheckpoint(DynamoUtils.safeGetLong(dynamoRecord, OWNER_SWITCHES_KEY));
result.setCheckpoint(
new ExtendedSequenceNumber(
DynamoUtils.safeGetString(dynamoRecord, CHECKPOINT_SEQUENCE_NUMBER_KEY),
DynamoUtils.safeGetLong(dynamoRecord, CHECKPOINT_SUBSEQUENCE_NUMBER_KEY))
);
result.setParentShardIds(DynamoUtils.safeGetSS(dynamoRecord, PARENT_SHARD_ID_KEY));
if (!Strings.isNullOrEmpty(DynamoUtils.safeGetString(dynamoRecord, PENDING_CHECKPOINT_SEQUENCE_KEY))) {
result.setPendingCheckpoint(
new ExtendedSequenceNumber(
DynamoUtils.safeGetString(dynamoRecord, PENDING_CHECKPOINT_SEQUENCE_KEY),
DynamoUtils.safeGetLong(dynamoRecord, PENDING_CHECKPOINT_SUBSEQUENCE_KEY))
);
}
return result;
}
@Override
public Map<String, AttributeValue> getDynamoHashKey(KinesisClientLease lease) {
return baseSerializer.getDynamoHashKey(lease);
}
@Override
public Map<String, AttributeValue> getDynamoHashKey(String shardId) {
return baseSerializer.getDynamoHashKey(shardId);
}
@Override
public Map<String, ExpectedAttributeValue> getDynamoLeaseCounterExpectation(KinesisClientLease lease) {
return baseSerializer.getDynamoLeaseCounterExpectation(lease);
}
@Override
public Map<String, ExpectedAttributeValue> getDynamoLeaseOwnerExpectation(KinesisClientLease lease) {
return baseSerializer.getDynamoLeaseOwnerExpectation(lease);
}
@Override
public Map<String, ExpectedAttributeValue> getDynamoNonexistantExpectation() {
return baseSerializer.getDynamoNonexistantExpectation();
}
@Override
public Map<String, AttributeValueUpdate> getDynamoLeaseCounterUpdate(KinesisClientLease lease) {
return baseSerializer.getDynamoLeaseCounterUpdate(lease);
}
@Override
public Map<String, AttributeValueUpdate> getDynamoTakeLeaseUpdate(KinesisClientLease lease, String newOwner) {
Map<String, AttributeValueUpdate> result = baseSerializer.getDynamoTakeLeaseUpdate(lease, newOwner);
String oldOwner = lease.getLeaseOwner();
if (oldOwner != null && !oldOwner.equals(newOwner)) {
result.put(OWNER_SWITCHES_KEY, new AttributeValueUpdate(DynamoUtils.createAttributeValue(1L),
AttributeAction.ADD));
}
return result;
}
@Override
public Map<String, AttributeValueUpdate> getDynamoEvictLeaseUpdate(KinesisClientLease lease) {
return baseSerializer.getDynamoEvictLeaseUpdate(lease);
}
@Override
public Map<String, AttributeValueUpdate> getDynamoUpdateLeaseUpdate(KinesisClientLease lease) {
Map<String, AttributeValueUpdate> result = baseSerializer.getDynamoUpdateLeaseUpdate(lease);
result.put(CHECKPOINT_SEQUENCE_NUMBER_KEY, new AttributeValueUpdate(DynamoUtils.createAttributeValue(lease.getCheckpoint().getSequenceNumber()),
AttributeAction.PUT));
result.put(CHECKPOINT_SUBSEQUENCE_NUMBER_KEY, new AttributeValueUpdate(DynamoUtils.createAttributeValue(lease.getCheckpoint().getSubSequenceNumber()),
AttributeAction.PUT));
result.put(OWNER_SWITCHES_KEY,
new AttributeValueUpdate(DynamoUtils.createAttributeValue(lease.getOwnerSwitchesSinceCheckpoint()),
AttributeAction.PUT));
if (lease.getPendingCheckpoint() != null && !lease.getPendingCheckpoint().getSequenceNumber().isEmpty()) {
result.put(PENDING_CHECKPOINT_SEQUENCE_KEY, new AttributeValueUpdate(DynamoUtils.createAttributeValue(lease.getPendingCheckpoint().getSequenceNumber()), AttributeAction.PUT));
result.put(PENDING_CHECKPOINT_SUBSEQUENCE_KEY, new AttributeValueUpdate(DynamoUtils.createAttributeValue(lease.getPendingCheckpoint().getSubSequenceNumber()), AttributeAction.PUT));
} else {
result.put(PENDING_CHECKPOINT_SEQUENCE_KEY, new AttributeValueUpdate().withAction(AttributeAction.DELETE));
result.put(PENDING_CHECKPOINT_SUBSEQUENCE_KEY, new AttributeValueUpdate().withAction(AttributeAction.DELETE));
}
return result;
}
@Override
public Collection<KeySchemaElement> getKeySchema() {
return baseSerializer.getKeySchema();
}
@Override
public Collection<AttributeDefinition> getAttributeDefinitions() {
return baseSerializer.getAttributeDefinitions();
}
}

View file

@ -1,153 +0,0 @@
/*
* Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Amazon Software License (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/asl/
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*/
package software.amazon.kinesis.leases;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import lombok.Getter;
import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
import software.amazon.kinesis.metrics.IMetricsFactory;
/**
* This class is used to coordinate/manage leases owned by this worker process and to get/set checkpoints.
*/
@Slf4j
public class KinesisClientLibLeaseCoordinator extends LeaseCoordinator<KinesisClientLease> {
private static final long DEFAULT_INITIAL_LEASE_TABLE_READ_CAPACITY = 10L;
private static final long DEFAULT_INITIAL_LEASE_TABLE_WRITE_CAPACITY = 10L;
@Getter
@Accessors(fluent = true)
private final LeaseManager<KinesisClientLease> leaseManager;
private long initialLeaseTableReadCapacity = DEFAULT_INITIAL_LEASE_TABLE_READ_CAPACITY;
private long initialLeaseTableWriteCapacity = DEFAULT_INITIAL_LEASE_TABLE_WRITE_CAPACITY;
public KinesisClientLibLeaseCoordinator(final LeaseManager<KinesisClientLease> leaseManager,
final String workerIdentifier,
final long leaseDurationMillis,
final long epsilonMillis,
final int maxLeasesForWorker,
final int maxLeasesToStealAtOneTime,
final int maxLeaseRenewerThreadCount,
final IMetricsFactory metricsFactory) {
super(leaseManager, workerIdentifier, leaseDurationMillis, epsilonMillis, maxLeasesForWorker,
maxLeasesToStealAtOneTime, maxLeaseRenewerThreadCount, metricsFactory);
this.leaseManager = leaseManager;
}
/**
* @param readCapacity The DynamoDB table used for tracking leases will be provisioned with the specified initial
* read capacity
* @return KinesisClientLibLeaseCoordinator
*/
public KinesisClientLibLeaseCoordinator initialLeaseTableReadCapacity(long readCapacity) {
if (readCapacity <= 0) {
throw new IllegalArgumentException("readCapacity should be >= 1");
}
this.initialLeaseTableReadCapacity = readCapacity;
return this;
}
/**
* @param writeCapacity The DynamoDB table used for tracking leases will be provisioned with the specified initial
* write capacity
* @return KinesisClientLibLeaseCoordinator
*/
public KinesisClientLibLeaseCoordinator initialLeaseTableWriteCapacity(long writeCapacity) {
if (writeCapacity <= 0) {
throw new IllegalArgumentException("writeCapacity should be >= 1");
}
this.initialLeaseTableWriteCapacity = writeCapacity;
return this;
}
/**
* @return Current shard/lease assignments
*/
public List<ShardInfo> getCurrentAssignments() {
Collection<KinesisClientLease> leases = getAssignments();
return convertLeasesToAssignments(leases);
}
public static List<ShardInfo> convertLeasesToAssignments(Collection<KinesisClientLease> leases) {
if (leases == null || leases.isEmpty()) {
return Collections.emptyList();
}
List<ShardInfo> assignments = new ArrayList<>(leases.size());
for (KinesisClientLease lease : leases) {
assignments.add(convertLeaseToAssignment(lease));
}
return assignments;
}
public static ShardInfo convertLeaseToAssignment(KinesisClientLease lease) {
Set<String> parentShardIds = lease.getParentShardIds();
return new ShardInfo(lease.getLeaseKey(), lease.getConcurrencyToken().toString(), parentShardIds,
lease.getCheckpoint());
}
/**
* Initialize the lease coordinator (create the lease table if needed).
* @throws DependencyException
* @throws ProvisionedThroughputException
*/
public void initialize() throws ProvisionedThroughputException, DependencyException, IllegalStateException {
final boolean newTableCreated =
leaseManager.createLeaseTableIfNotExists(initialLeaseTableReadCapacity, initialLeaseTableWriteCapacity);
if (newTableCreated) {
log.info("Created new lease table for coordinator with initial read capacity of {} and write capacity of {}.",
initialLeaseTableReadCapacity, initialLeaseTableWriteCapacity);
}
// Need to wait for table in active state.
final long secondsBetweenPolls = 10L;
final long timeoutSeconds = 600L;
final boolean isTableActive = leaseManager.waitUntilLeaseTableExists(secondsBetweenPolls, timeoutSeconds);
if (!isTableActive) {
throw new DependencyException(new IllegalStateException("Creating table timeout"));
}
}
/**
* Package access for testing.
*
* @throws DependencyException
* @throws InvalidStateException
*/
public void runLeaseTaker() throws DependencyException, InvalidStateException {
super.runTaker();
}
/**
* Package access for testing.
*
* @throws DependencyException
* @throws InvalidStateException
*/
public void runLeaseRenewer() throws DependencyException, InvalidStateException {
super.runRenewer();
}
}

View file

@ -36,7 +36,7 @@ import lombok.extern.slf4j.Slf4j;
*/ */
@RequiredArgsConstructor @RequiredArgsConstructor
@Slf4j @Slf4j
public class KinesisLeaseManagerProxy implements LeaseManagerProxy { public class KinesisShardDetector implements ShardDetector {
@NonNull @NonNull
private final AmazonKinesis amazonKinesis; private final AmazonKinesis amazonKinesis;
@NonNull @NonNull

View file

@ -14,10 +14,20 @@
*/ */
package software.amazon.kinesis.leases; package software.amazon.kinesis.leases;
import java.util.Collection;
import java.util.HashSet;
import java.util.Set;
import java.util.UUID; import java.util.UUID;
import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeUnit;
import com.amazonaws.util.json.Jackson; import com.amazonaws.util.CollectionUtils;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.NonNull;
import lombok.ToString;
import lombok.experimental.Accessors;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
/** /**
* This class contains data pertaining to a Lease. Distributed systems may use leases to partition work across a * This class contains data pertaining to a Lease. Distributed systems may use leases to partition work across a
@ -26,6 +36,11 @@ import com.amazonaws.util.json.Jackson;
* processing the corresponding unit of work, or until it fails. When the worker stops holding the lease, another worker will * processing the corresponding unit of work, or until it fails. When the worker stops holding the lease, another worker will
* take and hold the lease. * take and hold the lease.
*/ */
@NoArgsConstructor
@Getter
@Accessors(fluent = true)
@EqualsAndHashCode(exclude = {"concurrencyToken", "lastCounterIncrementNanos"})
@ToString
public class Lease { public class Lease {
/* /*
* See javadoc for System.nanoTime - summary: * See javadoc for System.nanoTime - summary:
@ -35,8 +50,17 @@ public class Lease {
*/ */
private static final long MAX_ABS_AGE_NANOS = TimeUnit.DAYS.toNanos(365); private static final long MAX_ABS_AGE_NANOS = TimeUnit.DAYS.toNanos(365);
/**
* @return leaseKey - identifies the unit of work associated with this lease.
*/
private String leaseKey; private String leaseKey;
/**
* @return current owner of the lease, may be null.
*/
private String leaseOwner; private String leaseOwner;
/**
* @return leaseCounter is incremented periodically by the holder of the lease. Used for optimistic locking.
*/
private Long leaseCounter = 0L; private Long leaseCounter = 0L;
/* /*
@ -50,12 +74,20 @@ public class Lease {
* deliberately not persisted in DynamoDB and excluded from hashCode and equals. * deliberately not persisted in DynamoDB and excluded from hashCode and equals.
*/ */
private Long lastCounterIncrementNanos; private Long lastCounterIncrementNanos;
/** /**
* Constructor. * @return most recently application-supplied checkpoint value. During fail over, the new worker will pick up after
* the old worker's last checkpoint.
*/ */
public Lease() { private ExtendedSequenceNumber checkpoint;
} /**
* @return pending checkpoint, possibly null.
*/
private ExtendedSequenceNumber pendingCheckpoint;
/**
* @return count of distinct lease holders between checkpoints.
*/
private Long ownerSwitchesSinceCheckpoint = 0L;
private Set<String> parentShardIds = new HashSet<>();
/** /**
* Copy constructor, used by clone(). * Copy constructor, used by clone().
@ -63,62 +95,46 @@ public class Lease {
* @param lease lease to copy * @param lease lease to copy
*/ */
protected Lease(Lease lease) { protected Lease(Lease lease) {
this(lease.getLeaseKey(), lease.getLeaseOwner(), lease.getLeaseCounter(), lease.getConcurrencyToken(), this(lease.leaseKey(), lease.leaseOwner(), lease.leaseCounter(), lease.concurrencyToken(),
lease.getLastCounterIncrementNanos()); lease.lastCounterIncrementNanos(), lease.checkpoint(), lease.pendingCheckpoint(),
lease.ownerSwitchesSinceCheckpoint(), lease.parentShardIds());
} }
protected Lease(String leaseKey, String leaseOwner, Long leaseCounter, UUID concurrencyToken, public Lease(final String leaseKey, final String leaseOwner, final Long leaseCounter,
Long lastCounterIncrementNanos) { final UUID concurrencyToken, final Long lastCounterIncrementNanos,
final ExtendedSequenceNumber checkpoint, final ExtendedSequenceNumber pendingCheckpoint,
final Long ownerSwitchesSinceCheckpoint, final Set<String> parentShardIds) {
this.leaseKey = leaseKey; this.leaseKey = leaseKey;
this.leaseOwner = leaseOwner; this.leaseOwner = leaseOwner;
this.leaseCounter = leaseCounter; this.leaseCounter = leaseCounter;
this.concurrencyToken = concurrencyToken; this.concurrencyToken = concurrencyToken;
this.lastCounterIncrementNanos = lastCounterIncrementNanos; this.lastCounterIncrementNanos = lastCounterIncrementNanos;
this.checkpoint = checkpoint;
this.pendingCheckpoint = pendingCheckpoint;
this.ownerSwitchesSinceCheckpoint = ownerSwitchesSinceCheckpoint;
if (!CollectionUtils.isNullOrEmpty(parentShardIds)) {
this.parentShardIds.addAll(parentShardIds);
}
}
/**
* @return shardIds that parent this lease. Used for resharding.
*/
public Set<String> parentShardIds() {
return new HashSet<>(parentShardIds);
} }
/** /**
* Updates this Lease's mutable, application-specific fields based on the passed-in lease object. Does not update * Updates this Lease's mutable, application-specific fields based on the passed-in lease object. Does not update
* fields that are internal to the leasing library (leaseKey, leaseOwner, leaseCounter). * fields that are internal to the leasing library (leaseKey, leaseOwner, leaseCounter).
* *
* @param other * @param lease
*/ */
public <T extends Lease> void update(T other) { public void update(final Lease lease) {
// The default implementation (no application-specific fields) has nothing to do. ownerSwitchesSinceCheckpoint(lease.ownerSwitchesSinceCheckpoint());
} checkpoint(lease.checkpoint);
pendingCheckpoint(lease.pendingCheckpoint);
/** parentShardIds(lease.parentShardIds);
* @return leaseKey - identifies the unit of work associated with this lease.
*/
public String getLeaseKey() {
return leaseKey;
}
/**
* @return leaseCounter is incremented periodically by the holder of the lease. Used for optimistic locking.
*/
public Long getLeaseCounter() {
return leaseCounter;
}
/**
* @return current owner of the lease, may be null.
*/
public String getLeaseOwner() {
return leaseOwner;
}
/**
* @return concurrency token
*/
public UUID getConcurrencyToken() {
return concurrencyToken;
}
/**
* @return last update in nanoseconds since the epoch
*/
public Long getLastCounterIncrementNanos() {
return lastCounterIncrementNanos;
} }
/** /**
@ -145,7 +161,7 @@ public class Lease {
* *
* @param lastCounterIncrementNanos last renewal in nanoseconds since the epoch * @param lastCounterIncrementNanos last renewal in nanoseconds since the epoch
*/ */
public void setLastCounterIncrementNanos(Long lastCounterIncrementNanos) { public void lastCounterIncrementNanos(Long lastCounterIncrementNanos) {
this.lastCounterIncrementNanos = lastCounterIncrementNanos; this.lastCounterIncrementNanos = lastCounterIncrementNanos;
} }
@ -154,8 +170,7 @@ public class Lease {
* *
* @param concurrencyToken may not be null * @param concurrencyToken may not be null
*/ */
public void setConcurrencyToken(UUID concurrencyToken) { public void concurrencyToken(@NonNull final UUID concurrencyToken) {
verifyNotNull(concurrencyToken, "concurencyToken cannot be null");
this.concurrencyToken = concurrencyToken; this.concurrencyToken = concurrencyToken;
} }
@ -164,12 +179,10 @@ public class Lease {
* *
* @param leaseKey may not be null. * @param leaseKey may not be null.
*/ */
public void setLeaseKey(String leaseKey) { public void leaseKey(@NonNull final String leaseKey) {
if (this.leaseKey != null) { if (this.leaseKey != null) {
throw new IllegalArgumentException("LeaseKey is immutable once set"); throw new IllegalArgumentException("LeaseKey is immutable once set");
} }
verifyNotNull(leaseKey, "LeaseKey cannot be set to null");
this.leaseKey = leaseKey; this.leaseKey = leaseKey;
} }
@ -178,77 +191,62 @@ public class Lease {
* *
* @param leaseCounter may not be null * @param leaseCounter may not be null
*/ */
public void setLeaseCounter(Long leaseCounter) { public void leaseCounter(@NonNull final Long leaseCounter) {
verifyNotNull(leaseCounter, "leaseCounter must not be null");
this.leaseCounter = leaseCounter; this.leaseCounter = leaseCounter;
} }
/**
* Sets checkpoint.
*
* @param checkpoint may not be null
*/
public void checkpoint(@NonNull final ExtendedSequenceNumber checkpoint) {
this.checkpoint = checkpoint;
}
/**
* Sets pending checkpoint.
*
* @param pendingCheckpoint can be null
*/
public void pendingCheckpoint(ExtendedSequenceNumber pendingCheckpoint) {
this.pendingCheckpoint = pendingCheckpoint;
}
/**
* Sets ownerSwitchesSinceCheckpoint.
*
* @param ownerSwitchesSinceCheckpoint may not be null
*/
public void ownerSwitchesSinceCheckpoint(@NonNull final Long ownerSwitchesSinceCheckpoint) {
this.ownerSwitchesSinceCheckpoint = ownerSwitchesSinceCheckpoint;
}
/**
* Sets parentShardIds.
*
* @param parentShardIds may not be null
*/
public void parentShardIds(@NonNull final Collection<String> parentShardIds) {
this.parentShardIds.clear();
this.parentShardIds.addAll(parentShardIds);
}
/** /**
* Sets leaseOwner. * Sets leaseOwner.
* *
* @param leaseOwner may be null. * @param leaseOwner may be null.
*/ */
public void setLeaseOwner(String leaseOwner) { public void leaseOwner(String leaseOwner) {
this.leaseOwner = leaseOwner; this.leaseOwner = leaseOwner;
} }
@Override
public int hashCode() {
final int prime = 31;
int result = 1;
result = prime * result + ((leaseCounter == null) ? 0 : leaseCounter.hashCode());
result = prime * result + ((leaseOwner == null) ? 0 : leaseOwner.hashCode());
result = prime * result + ((leaseKey == null) ? 0 : leaseKey.hashCode());
return result;
}
@Override
public boolean equals(Object obj) {
if (this == obj)
return true;
if (obj == null)
return false;
if (getClass() != obj.getClass())
return false;
Lease other = (Lease) obj;
if (leaseCounter == null) {
if (other.leaseCounter != null)
return false;
} else if (!leaseCounter.equals(other.leaseCounter))
return false;
if (leaseOwner == null) {
if (other.leaseOwner != null)
return false;
} else if (!leaseOwner.equals(other.leaseOwner))
return false;
if (leaseKey == null) {
if (other.leaseKey != null)
return false;
} else if (!leaseKey.equals(other.leaseKey))
return false;
return true;
}
@Override
public String toString() {
return Jackson.toJsonPrettyString(this);
}
/** /**
* Returns a deep copy of this object. Type-unsafe - there aren't good mechanisms for copy-constructing generics. * Returns a deep copy of this object. Type-unsafe - there aren't good mechanisms for copy-constructing generics.
* *
* @return A deep copy of this object. * @return A deep copy of this object.
*/ */
@SuppressWarnings("unchecked") public Lease copy() {
public <T extends Lease> T copy() { return new Lease(this);
return (T) new Lease(this);
} }
private void verifyNotNull(Object object, String message) {
if (object == null) {
throw new IllegalArgumentException(message);
}
}
} }

View file

@ -12,175 +12,28 @@
* express or implied. See the License for the specific language governing * express or implied. See the License for the specific language governing
* permissions and limitations under the License. * permissions and limitations under the License.
*/ */
package software.amazon.kinesis.leases; package software.amazon.kinesis.leases;
import java.util.Collection; import java.util.Collection;
import java.util.Map; import java.util.List;
import java.util.UUID; import java.util.UUID;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedTransferQueue;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import com.google.common.util.concurrent.ThreadFactoryBuilder; import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseCoordinator;
import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.LeasingException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException; import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
import software.amazon.kinesis.metrics.IMetricsFactory;
import software.amazon.kinesis.metrics.IMetricsScope;
import software.amazon.kinesis.metrics.LogMetricsFactory;
import software.amazon.kinesis.metrics.MetricsHelper;
import software.amazon.kinesis.metrics.MetricsLevel;
/** /**
* LeaseCoordinator abstracts away LeaseTaker and LeaseRenewer from the application code that's using leasing. It owns
* the scheduling of the two previously mentioned components as well as informing LeaseRenewer when LeaseTaker takes new
* leases.
* *
*/ */
@Slf4j public interface LeaseCoordinator {
public class LeaseCoordinator<T extends Lease> {
/*
* Name of the dimension used when setting worker identifier on IMetricsScopes. Exposed so that users of this class
* can easily create InterceptingMetricsFactories that rename this dimension to suit the destination metrics system.
*/
public static final String WORKER_IDENTIFIER_METRIC = "WorkerIdentifier";
// Time to wait for in-flight Runnables to finish when calling .stop();
private static final long STOP_WAIT_TIME_MILLIS = 2000L;
private static final int DEFAULT_MAX_LEASES_FOR_WORKER = Integer.MAX_VALUE;
private static final int DEFAULT_MAX_LEASES_TO_STEAL_AT_ONE_TIME = 1;
private static final ThreadFactory LEASE_COORDINATOR_THREAD_FACTORY = new ThreadFactoryBuilder()
.setNameFormat("LeaseCoordinator-%04d").setDaemon(true).build();
private static final ThreadFactory LEASE_RENEWAL_THREAD_FACTORY = new ThreadFactoryBuilder()
.setNameFormat("LeaseRenewer-%04d").setDaemon(true).build();
private final LeaseRenewer<T> leaseRenewer;
private final LeaseTaker<T> leaseTaker;
private final long renewerIntervalMillis;
private final long takerIntervalMillis;
private final Object shutdownLock = new Object();
protected final IMetricsFactory metricsFactory;
private ScheduledExecutorService leaseCoordinatorThreadPool;
private final ExecutorService leaseRenewalThreadpool;
private volatile boolean running = false;
private ScheduledFuture<?> takerFuture;
/** /**
* Constructor. * Initialize the lease coordinator (create the lease table if needed).
* * @throws DependencyException
* @param leaseManager LeaseManager instance to use * @throws ProvisionedThroughputException
* @param workerIdentifier Identifies the worker (e.g. useful to track lease ownership)
* @param leaseDurationMillis Duration of a lease
* @param epsilonMillis Allow for some variance when calculating lease expirations
*/ */
public LeaseCoordinator(LeaseManager<T> leaseManager, void initialize() throws ProvisionedThroughputException, DependencyException, IllegalStateException;
String workerIdentifier,
long leaseDurationMillis,
long epsilonMillis) {
this(leaseManager, workerIdentifier, leaseDurationMillis, epsilonMillis, new LogMetricsFactory());
}
/**
* Constructor.
*
* @param leaseManager LeaseManager instance to use
* @param workerIdentifier Identifies the worker (e.g. useful to track lease ownership)
* @param leaseDurationMillis Duration of a lease
* @param epsilonMillis Allow for some variance when calculating lease expirations
* @param metricsFactory Used to publish metrics about lease operations
*/
public LeaseCoordinator(LeaseManager<T> leaseManager,
String workerIdentifier,
long leaseDurationMillis,
long epsilonMillis,
IMetricsFactory metricsFactory) {
this(leaseManager, workerIdentifier, leaseDurationMillis, epsilonMillis,
DEFAULT_MAX_LEASES_FOR_WORKER, DEFAULT_MAX_LEASES_TO_STEAL_AT_ONE_TIME,
LeaseManagementConfig.DEFAULT_MAX_LEASE_RENEWAL_THREADS, metricsFactory);
}
/**
* Constructor.
*
* @param leaseManager LeaseManager instance to use
* @param workerIdentifier Identifies the worker (e.g. useful to track lease ownership)
* @param leaseDurationMillis Duration of a lease
* @param epsilonMillis Allow for some variance when calculating lease expirations
* @param maxLeasesForWorker Max leases this Worker can handle at a time
* @param maxLeasesToStealAtOneTime Steal up to these many leases at a time (for load balancing)
* @param metricsFactory Used to publish metrics about lease operations
*/
public LeaseCoordinator(LeaseManager<T> leaseManager,
String workerIdentifier,
long leaseDurationMillis,
long epsilonMillis,
int maxLeasesForWorker,
int maxLeasesToStealAtOneTime,
int maxLeaseRenewerThreadCount,
IMetricsFactory metricsFactory) {
this.leaseRenewalThreadpool = getLeaseRenewalExecutorService(maxLeaseRenewerThreadCount);
this.leaseTaker = new DynamoDBLeaseTaker<T>(leaseManager, workerIdentifier, leaseDurationMillis)
.withMaxLeasesForWorker(maxLeasesForWorker)
.withMaxLeasesToStealAtOneTime(maxLeasesToStealAtOneTime);
this.leaseRenewer = new DynamoDBLeaseRenewer<T>(
leaseManager, workerIdentifier, leaseDurationMillis, leaseRenewalThreadpool);
this.renewerIntervalMillis = leaseDurationMillis / 3 - epsilonMillis;
this.takerIntervalMillis = (leaseDurationMillis + epsilonMillis) * 2;
this.metricsFactory = metricsFactory;
log.info("With failover time {} ms and epsilon {} ms, LeaseCoordinator will renew leases every {} ms, take"
+ "leases every {} ms, process maximum of {} leases and steal {} lease(s) at a time.",
leaseDurationMillis,
epsilonMillis,
renewerIntervalMillis,
takerIntervalMillis,
maxLeasesForWorker,
maxLeasesToStealAtOneTime);
}
private class TakerRunnable implements Runnable {
@Override
public void run() {
try {
runTaker();
} catch (LeasingException e) {
log.error("LeasingException encountered in lease taking thread", e);
} catch (Throwable t) {
log.error("Throwable encountered in lease taking thread", t);
}
}
}
private class RenewerRunnable implements Runnable {
@Override
public void run() {
try {
runRenewer();
} catch (LeasingException e) {
log.error("LeasingException encountered in lease renewing thread", e);
} catch (Throwable t) {
log.error("Throwable encountered in lease renewing thread", t);
}
}
}
/** /**
* Start background LeaseHolder and LeaseTaker threads. * Start background LeaseHolder and LeaseTaker threads.
@ -188,24 +41,7 @@ public class LeaseCoordinator<T extends Lease> {
* @throws InvalidStateException If the lease table doesn't exist * @throws InvalidStateException If the lease table doesn't exist
* @throws DependencyException If we encountered exception taking to DynamoDB * @throws DependencyException If we encountered exception taking to DynamoDB
*/ */
public void start() throws DependencyException, InvalidStateException, ProvisionedThroughputException { void start() throws DependencyException, InvalidStateException, ProvisionedThroughputException;
leaseRenewer.initialize();
// 2 because we know we'll have at most 2 concurrent tasks at a time.
leaseCoordinatorThreadPool = Executors.newScheduledThreadPool(2, LEASE_COORDINATOR_THREAD_FACTORY);
// Taker runs with fixed DELAY because we want it to run slower in the event of performance degredation.
takerFuture = leaseCoordinatorThreadPool.scheduleWithFixedDelay(new TakerRunnable(),
0L,
takerIntervalMillis,
TimeUnit.MILLISECONDS);
// Renewer runs at fixed INTERVAL because we want it to run at the same rate in the event of degredation.
leaseCoordinatorThreadPool.scheduleAtFixedRate(new RenewerRunnable(),
0L,
renewerIntervalMillis,
TimeUnit.MILLISECONDS);
running = true;
}
/** /**
* Runs a single iteration of the lease taker - used by integration tests. * Runs a single iteration of the lease taker - used by integration tests.
@ -213,28 +49,7 @@ public class LeaseCoordinator<T extends Lease> {
* @throws InvalidStateException * @throws InvalidStateException
* @throws DependencyException * @throws DependencyException
*/ */
protected void runTaker() throws DependencyException, InvalidStateException { void runLeaseTaker() throws DependencyException, InvalidStateException;
IMetricsScope scope = MetricsHelper.startScope(metricsFactory, "TakeLeases");
long startTime = System.currentTimeMillis();
boolean success = false;
try {
Map<String, T> takenLeases = leaseTaker.takeLeases();
// Only add taken leases to renewer if coordinator is still running.
synchronized (shutdownLock) {
if (running) {
leaseRenewer.addLeasesToRenew(takenLeases.values());
}
}
success = true;
} finally {
scope.addDimension(WORKER_IDENTIFIER_METRIC, getWorkerIdentifier());
MetricsHelper.addSuccessAndLatency(startTime, success, MetricsLevel.SUMMARY);
MetricsHelper.endScope();
}
}
/** /**
* Runs a single iteration of the lease renewer - used by integration tests. * Runs a single iteration of the lease renewer - used by integration tests.
@ -242,108 +57,40 @@ public class LeaseCoordinator<T extends Lease> {
* @throws InvalidStateException * @throws InvalidStateException
* @throws DependencyException * @throws DependencyException
*/ */
protected void runRenewer() throws DependencyException, InvalidStateException { void runLeaseRenewer() throws DependencyException, InvalidStateException;
IMetricsScope scope = MetricsHelper.startScope(metricsFactory, "RenewAllLeases");
long startTime = System.currentTimeMillis();
boolean success = false;
try { /**
leaseRenewer.renewLeases(); * @return true if this LeaseCoordinator is running
success = true; */
} finally { boolean isRunning();
scope.addDimension(WORKER_IDENTIFIER_METRIC, getWorkerIdentifier());
MetricsHelper.addSuccessAndLatency(startTime, success, MetricsLevel.SUMMARY); /**
MetricsHelper.endScope(); * @return workerIdentifier
} */
} String workerIdentifier();
/**
* @return {@link LeaseRefresher}
*/
LeaseRefresher leaseRefresher();
/** /**
* @return currently held leases * @return currently held leases
*/ */
public Collection<T> getAssignments() { Collection<Lease> getAssignments();
return leaseRenewer.getCurrentlyHeldLeases().values();
}
/** /**
* @param leaseKey lease key to fetch currently held lease for * @param leaseKey lease key to fetch currently held lease for
* *
* @return deep copy of currently held Lease for given key, or null if we don't hold the lease for that key * @return deep copy of currently held Lease for given key, or null if we don't hold the lease for that key
*/ */
public T getCurrentlyHeldLease(String leaseKey) { Lease getCurrentlyHeldLease(String leaseKey);
return leaseRenewer.getCurrentlyHeldLease(leaseKey);
}
/**
* @return workerIdentifier
*/
public String getWorkerIdentifier() {
return leaseTaker.getWorkerIdentifier();
}
/**
* Stops background threads and waits for {@link #STOP_WAIT_TIME_MILLIS} for all background tasks to complete.
* If tasks are not completed after this time, method will shutdown thread pool forcefully and return.
*/
public void stop() {
if (leaseCoordinatorThreadPool != null) {
leaseCoordinatorThreadPool.shutdown();
try {
if (leaseCoordinatorThreadPool.awaitTermination(STOP_WAIT_TIME_MILLIS, TimeUnit.MILLISECONDS)) {
log.info("Worker {} has successfully stopped lease-tracking threads",
leaseTaker.getWorkerIdentifier());
} else {
leaseCoordinatorThreadPool.shutdownNow();
log.info("Worker {} stopped lease-tracking threads {} ms after stop",
leaseTaker.getWorkerIdentifier(),
STOP_WAIT_TIME_MILLIS);
}
} catch (InterruptedException e) {
log.debug("Encountered InterruptedException when awaiting threadpool termination");
}
} else {
log.debug("Threadpool was null, no need to shutdown/terminate threadpool.");
}
leaseRenewalThreadpool.shutdownNow();
synchronized (shutdownLock) {
leaseRenewer.clearCurrentlyHeldLeases();
running = false;
}
}
/**
* Requests the cancellation of the lease taker.
*/
public void stopLeaseTaker() {
takerFuture.cancel(false);
}
/**
* Requests that renewals for the given lease are stopped.
*
* @param lease the lease to stop renewing.
*/
public void dropLease(T lease) {
synchronized (shutdownLock) {
if (lease != null) {
leaseRenewer.dropLease(lease);
}
}
}
/**
* @return true if this LeaseCoordinator is running
*/
public boolean isRunning() {
return running;
}
/** /**
* Updates application-specific lease values in DynamoDB. * Updates application-specific lease values in DynamoDB.
* *
* @param lease lease object containing updated values * @param lease lease object containing updated values
* @param concurrencyToken obtained by calling Lease.getConcurrencyToken for a currently held lease * @param concurrencyToken obtained by calling Lease.concurrencyToken for a currently held lease
* *
* @return true if update succeeded, false otherwise * @return true if update succeeded, false otherwise
* *
@ -351,20 +98,43 @@ public class LeaseCoordinator<T extends Lease> {
* @throws ProvisionedThroughputException if DynamoDB update fails due to lack of capacity * @throws ProvisionedThroughputException if DynamoDB update fails due to lack of capacity
* @throws DependencyException if DynamoDB update fails in an unexpected way * @throws DependencyException if DynamoDB update fails in an unexpected way
*/ */
public boolean updateLease(T lease, UUID concurrencyToken) boolean updateLease(Lease lease, UUID concurrencyToken)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException;
return leaseRenewer.updateLease(lease, concurrencyToken);
}
/** /**
* Returns executor service that should be used for lease renewal. * Requests the cancellation of the lease taker.
* @param maximumPoolSize Maximum allowed thread pool size
* @return Executor service that should be used for lease renewal.
*/ */
private static ExecutorService getLeaseRenewalExecutorService(int maximumPoolSize) { void stopLeaseTaker();
int coreLeaseCount = Math.max(maximumPoolSize / 4, 2);
return new ThreadPoolExecutor(coreLeaseCount, maximumPoolSize, 60, TimeUnit.SECONDS, /**
new LinkedTransferQueue<Runnable>(), LEASE_RENEWAL_THREAD_FACTORY); * Requests that renewals for the given lease are stopped.
} *
* @param lease the lease to stop renewing.
*/
void dropLease(Lease lease);
/**
* Stops background threads and waits for specific amount of time for all background tasks to complete.
* If tasks are not completed after this time, method will shutdown thread pool forcefully and return.
*/
void stop();
/**
* @return Current shard/lease assignments
*/
List<ShardInfo> getCurrentAssignments();
/**
* @param writeCapacity The DynamoDB table used for tracking leases will be provisioned with the specified initial
* write capacity
* @return LeaseCoordinator
*/
DynamoDBLeaseCoordinator initialLeaseTableWriteCapacity(long writeCapacity);
/**
* @param readCapacity The DynamoDB table used for tracking leases will be provisioned with the specified initial
* read capacity
* @return LeaseCoordinator
*/
DynamoDBLeaseCoordinator initialLeaseTableReadCapacity(long readCapacity);
} }

View file

@ -30,6 +30,7 @@ import com.google.common.util.concurrent.ThreadFactoryBuilder;
import lombok.Data; import lombok.Data;
import lombok.NonNull; import lombok.NonNull;
import lombok.experimental.Accessors; import lombok.experimental.Accessors;
import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseManagementFactory;
import software.amazon.kinesis.metrics.IMetricsFactory; import software.amazon.kinesis.metrics.IMetricsFactory;
import software.amazon.kinesis.metrics.NullMetricsFactory; import software.amazon.kinesis.metrics.NullMetricsFactory;

View file

@ -15,6 +15,8 @@
package software.amazon.kinesis.leases; package software.amazon.kinesis.leases;
import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseRefresher;
/** /**
* *
*/ */
@ -23,9 +25,7 @@ public interface LeaseManagementFactory {
ShardSyncTaskManager createShardSyncTaskManager(); ShardSyncTaskManager createShardSyncTaskManager();
DynamoDBLeaseManager<KinesisClientLease> createLeaseManager(); DynamoDBLeaseRefresher createLeaseRefresher();
KinesisClientLibLeaseCoordinator createKinesisClientLibLeaseCoordinator(); ShardDetector createShardDetector();
LeaseManagerProxy createLeaseManagerProxy();
} }

View file

@ -19,14 +19,12 @@ import java.util.List;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException; import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
import software.amazon.kinesis.leases.Lease; import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
/** /**
* Supports basic CRUD operations for Leases. * Supports basic CRUD operations for Leases.
*
* @param <T> Lease subclass, possibly Lease itself.
*/ */
public interface LeaseManager<T extends Lease> { public interface LeaseRefresher {
/** /**
* Creates the table that will store leases. Succeeds if table already exists. * Creates the table that will store leases. Succeeds if table already exists.
@ -40,7 +38,7 @@ public interface LeaseManager<T extends Lease> {
* restrictions. * restrictions.
* @throws DependencyException if DynamoDB createTable fails in an unexpected way * @throws DependencyException if DynamoDB createTable fails in an unexpected way
*/ */
public boolean createLeaseTableIfNotExists(Long readCapacity, Long writeCapacity) boolean createLeaseTableIfNotExists(Long readCapacity, Long writeCapacity)
throws ProvisionedThroughputException, DependencyException; throws ProvisionedThroughputException, DependencyException;
/** /**
@ -48,7 +46,7 @@ public interface LeaseManager<T extends Lease> {
* *
* @throws DependencyException if DynamoDB describeTable fails in an unexpected way * @throws DependencyException if DynamoDB describeTable fails in an unexpected way
*/ */
public boolean leaseTableExists() throws DependencyException; boolean leaseTableExists() throws DependencyException;
/** /**
* Blocks until the lease table exists by polling leaseTableExists. * Blocks until the lease table exists by polling leaseTableExists.
@ -60,7 +58,7 @@ public interface LeaseManager<T extends Lease> {
* *
* @throws DependencyException if DynamoDB describeTable fails in an unexpected way * @throws DependencyException if DynamoDB describeTable fails in an unexpected way
*/ */
public boolean waitUntilLeaseTableExists(long secondsBetweenPolls, long timeoutSeconds) throws DependencyException; boolean waitUntilLeaseTableExists(long secondsBetweenPolls, long timeoutSeconds) throws DependencyException;
/** /**
* List all objects in table synchronously. * List all objects in table synchronously.
@ -71,7 +69,7 @@ public interface LeaseManager<T extends Lease> {
* *
* @return list of leases * @return list of leases
*/ */
public List<T> listLeases() throws DependencyException, InvalidStateException, ProvisionedThroughputException; List<Lease> listLeases() throws DependencyException, InvalidStateException, ProvisionedThroughputException;
/** /**
* Create a new lease. Conditional on a lease not already existing with this shardId. * Create a new lease. Conditional on a lease not already existing with this shardId.
@ -84,7 +82,7 @@ public interface LeaseManager<T extends Lease> {
* @throws InvalidStateException if lease table does not exist * @throws InvalidStateException if lease table does not exist
* @throws ProvisionedThroughputException if DynamoDB put fails due to lack of capacity * @throws ProvisionedThroughputException if DynamoDB put fails due to lack of capacity
*/ */
public boolean createLeaseIfNotExists(T lease) boolean createLeaseIfNotExists(Lease lease)
throws DependencyException, InvalidStateException, ProvisionedThroughputException; throws DependencyException, InvalidStateException, ProvisionedThroughputException;
/** /**
@ -96,7 +94,7 @@ public interface LeaseManager<T extends Lease> {
* *
* @return lease for the specified shardId, or null if one doesn't exist * @return lease for the specified shardId, or null if one doesn't exist
*/ */
public T getLease(String shardId) throws DependencyException, InvalidStateException, ProvisionedThroughputException; Lease getLease(String shardId) throws DependencyException, InvalidStateException, ProvisionedThroughputException;
/** /**
* Renew a lease by incrementing the lease counter. Conditional on the leaseCounter in DynamoDB matching the leaseCounter * Renew a lease by incrementing the lease counter. Conditional on the leaseCounter in DynamoDB matching the leaseCounter
@ -110,7 +108,7 @@ public interface LeaseManager<T extends Lease> {
* @throws ProvisionedThroughputException if DynamoDB update fails due to lack of capacity * @throws ProvisionedThroughputException if DynamoDB update fails due to lack of capacity
* @throws DependencyException if DynamoDB update fails in an unexpected way * @throws DependencyException if DynamoDB update fails in an unexpected way
*/ */
public boolean renewLease(T lease) boolean renewLease(Lease lease)
throws DependencyException, InvalidStateException, ProvisionedThroughputException; throws DependencyException, InvalidStateException, ProvisionedThroughputException;
/** /**
@ -127,7 +125,7 @@ public interface LeaseManager<T extends Lease> {
* @throws ProvisionedThroughputException if DynamoDB update fails due to lack of capacity * @throws ProvisionedThroughputException if DynamoDB update fails due to lack of capacity
* @throws DependencyException if DynamoDB update fails in an unexpected way * @throws DependencyException if DynamoDB update fails in an unexpected way
*/ */
public boolean takeLease(T lease, String owner) boolean takeLease(Lease lease, String owner)
throws DependencyException, InvalidStateException, ProvisionedThroughputException; throws DependencyException, InvalidStateException, ProvisionedThroughputException;
/** /**
@ -142,7 +140,7 @@ public interface LeaseManager<T extends Lease> {
* @throws ProvisionedThroughputException if DynamoDB update fails due to lack of capacity * @throws ProvisionedThroughputException if DynamoDB update fails due to lack of capacity
* @throws DependencyException if DynamoDB update fails in an unexpected way * @throws DependencyException if DynamoDB update fails in an unexpected way
*/ */
public boolean evictLease(T lease) boolean evictLease(Lease lease)
throws DependencyException, InvalidStateException, ProvisionedThroughputException; throws DependencyException, InvalidStateException, ProvisionedThroughputException;
/** /**
@ -154,7 +152,7 @@ public interface LeaseManager<T extends Lease> {
* @throws ProvisionedThroughputException if DynamoDB delete fails due to lack of capacity * @throws ProvisionedThroughputException if DynamoDB delete fails due to lack of capacity
* @throws DependencyException if DynamoDB delete fails in an unexpected way * @throws DependencyException if DynamoDB delete fails in an unexpected way
*/ */
public void deleteLease(T lease) throws DependencyException, InvalidStateException, ProvisionedThroughputException; void deleteLease(Lease lease) throws DependencyException, InvalidStateException, ProvisionedThroughputException;
/** /**
* Delete all leases from DynamoDB. Useful for tools/utils and testing. * Delete all leases from DynamoDB. Useful for tools/utils and testing.
@ -163,7 +161,7 @@ public interface LeaseManager<T extends Lease> {
* @throws ProvisionedThroughputException if DynamoDB scan or delete fail due to lack of capacity * @throws ProvisionedThroughputException if DynamoDB scan or delete fail due to lack of capacity
* @throws DependencyException if DynamoDB scan or delete fail in an unexpected way * @throws DependencyException if DynamoDB scan or delete fail in an unexpected way
*/ */
public void deleteAll() throws DependencyException, InvalidStateException, ProvisionedThroughputException; void deleteAll() throws DependencyException, InvalidStateException, ProvisionedThroughputException;
/** /**
* Update application-specific fields of the given lease in DynamoDB. Does not update fields managed by the leasing * Update application-specific fields of the given lease in DynamoDB. Does not update fields managed by the leasing
@ -177,7 +175,7 @@ public interface LeaseManager<T extends Lease> {
* @throws ProvisionedThroughputException if DynamoDB update fails due to lack of capacity * @throws ProvisionedThroughputException if DynamoDB update fails due to lack of capacity
* @throws DependencyException if DynamoDB update fails in an unexpected way * @throws DependencyException if DynamoDB update fails in an unexpected way
*/ */
public boolean updateLease(T lease) boolean updateLease(Lease lease)
throws DependencyException, InvalidStateException, ProvisionedThroughputException; throws DependencyException, InvalidStateException, ProvisionedThroughputException;
/** /**
@ -189,6 +187,20 @@ public interface LeaseManager<T extends Lease> {
* @throws InvalidStateException if lease table does not exist * @throws InvalidStateException if lease table does not exist
* @throws ProvisionedThroughputException if DynamoDB scan fails due to lack of capacity * @throws ProvisionedThroughputException if DynamoDB scan fails due to lack of capacity
*/ */
public boolean isLeaseTableEmpty() throws DependencyException, InvalidStateException, ProvisionedThroughputException; boolean isLeaseTableEmpty() throws DependencyException, InvalidStateException, ProvisionedThroughputException;
/**
* Gets the current checkpoint of the shard. This is useful in the resharding use case
* where we will wait for the parent shard to complete before starting on the records from a child shard.
*
* @param shardId Checkpoint of this shard will be returned
* @return Checkpoint of this shard, or null if the shard record doesn't exist.
*
* @throws ProvisionedThroughputException if DynamoDB update fails due to lack of capacity
* @throws InvalidStateException if lease table does not exist
* @throws DependencyException if DynamoDB update fails in an unexpected way
*/
ExtendedSequenceNumber getCheckpoint(String shardId)
throws ProvisionedThroughputException, InvalidStateException, DependencyException;
} }

View file

@ -21,22 +21,21 @@ import java.util.UUID;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException; import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
import software.amazon.kinesis.leases.Lease;
/** /**
* ILeaseRenewer objects are used by LeaseCoordinator to renew leases held by the LeaseCoordinator. Each * ILeaseRenewer objects are used by LeaseCoordinator to renew leases held by the LeaseCoordinator. Each
* LeaseCoordinator instance corresponds to one worker, and uses exactly one ILeaseRenewer to manage lease renewal for * LeaseCoordinator instance corresponds to one worker, and uses exactly one ILeaseRenewer to manage lease renewal for
* that worker. * that worker.
*/ */
public interface LeaseRenewer<T extends Lease> { public interface LeaseRenewer {
/** /**
* Bootstrap initial set of leases from the LeaseManager (e.g. upon process restart, pick up leases we own) * Bootstrap initial set of leases from the {@link LeaseRefresher} (e.g. upon process restart, pick up leases we own)
* @throws DependencyException on unexpected DynamoDB failures * @throws DependencyException on unexpected DynamoDB failures
* @throws InvalidStateException if lease table doesn't exist * @throws InvalidStateException if lease table doesn't exist
* @throws ProvisionedThroughputException if DynamoDB reads fail due to insufficient capacity * @throws ProvisionedThroughputException if DynamoDB reads fail due to insufficient capacity
*/ */
public void initialize() throws DependencyException, InvalidStateException, ProvisionedThroughputException; void initialize() throws DependencyException, InvalidStateException, ProvisionedThroughputException;
/** /**
* Attempt to renew all currently held leases. * Attempt to renew all currently held leases.
@ -44,21 +43,21 @@ public interface LeaseRenewer<T extends Lease> {
* @throws DependencyException on unexpected DynamoDB failures * @throws DependencyException on unexpected DynamoDB failures
* @throws InvalidStateException if lease table does not exist * @throws InvalidStateException if lease table does not exist
*/ */
public void renewLeases() throws DependencyException, InvalidStateException; void renewLeases() throws DependencyException, InvalidStateException;
/** /**
* @return currently held leases. Key is shardId, value is corresponding Lease object. A lease is currently held if * @return currently held leases. Key is shardId, value is corresponding Lease object. A lease is currently held if
* we successfully renewed it on the last run of renewLeases(). Lease objects returned are deep copies - * we successfully renewed it on the last run of renewLeases(). Lease objects returned are deep copies -
* their lease counters will not tick. * their lease counters will not tick.
*/ */
public Map<String, T> getCurrentlyHeldLeases(); Map<String, Lease> getCurrentlyHeldLeases();
/** /**
* @param leaseKey key of the lease to retrieve * @param leaseKey key of the lease to retrieve
* *
* @return a deep copy of a currently held lease, or null if we don't hold the lease * @return a deep copy of a currently held lease, or null if we don't hold the lease
*/ */
public T getCurrentlyHeldLease(String leaseKey); Lease getCurrentlyHeldLease(String leaseKey);
/** /**
* Adds leases to this LeaseRenewer's set of currently held leases. Leases must have lastRenewalNanos set to the * Adds leases to this LeaseRenewer's set of currently held leases. Leases must have lastRenewalNanos set to the
@ -66,19 +65,19 @@ public interface LeaseRenewer<T extends Lease> {
* *
* @param newLeases new leases. * @param newLeases new leases.
*/ */
public void addLeasesToRenew(Collection<T> newLeases); void addLeasesToRenew(Collection<Lease> newLeases);
/** /**
* Clears this LeaseRenewer's set of currently held leases. * Clears this LeaseRenewer's set of currently held leases.
*/ */
public void clearCurrentlyHeldLeases(); void clearCurrentlyHeldLeases();
/** /**
* Stops the lease renewer from continunig to maintain the given lease. * Stops the lease renewer from continunig to maintain the given lease.
* *
* @param lease the lease to drop. * @param lease the lease to drop.
*/ */
void dropLease(T lease); void dropLease(Lease lease);
/** /**
* Update application-specific fields in a currently held lease. Cannot be used to update internal fields such as * Update application-specific fields in a currently held lease. Cannot be used to update internal fields such as
@ -86,7 +85,7 @@ public interface LeaseRenewer<T extends Lease> {
* the concurrency token on the internal authoritative copy of the lease (ie, if we lost and re-acquired the lease). * the concurrency token on the internal authoritative copy of the lease (ie, if we lost and re-acquired the lease).
* *
* @param lease lease object containing updated data * @param lease lease object containing updated data
* @param concurrencyToken obtained by calling Lease.getConcurrencyToken for a currently held lease * @param concurrencyToken obtained by calling Lease.concurrencyToken for a currently held lease
* *
* @return true if update succeeds, false otherwise * @return true if update succeeds, false otherwise
* *
@ -94,7 +93,7 @@ public interface LeaseRenewer<T extends Lease> {
* @throws ProvisionedThroughputException if DynamoDB update fails due to lack of capacity * @throws ProvisionedThroughputException if DynamoDB update fails due to lack of capacity
* @throws DependencyException if DynamoDB update fails in an unexpected way * @throws DependencyException if DynamoDB update fails in an unexpected way
*/ */
boolean updateLease(T lease, UUID concurrencyToken) boolean updateLease(Lease lease, UUID concurrencyToken)
throws DependencyException, InvalidStateException, ProvisionedThroughputException; throws DependencyException, InvalidStateException, ProvisionedThroughputException;
} }

View file

@ -26,10 +26,8 @@ import software.amazon.kinesis.leases.Lease;
/** /**
* Utility class that manages the mapping of Lease objects/operations to records in DynamoDB. * Utility class that manages the mapping of Lease objects/operations to records in DynamoDB.
*
* @param <T> Lease subclass, possibly Lease itself
*/ */
public interface LeaseSerializer<T extends Lease> { public interface LeaseSerializer {
/** /**
* Construct a DynamoDB record out of a Lease object * Construct a DynamoDB record out of a Lease object
@ -37,7 +35,7 @@ public interface LeaseSerializer<T extends Lease> {
* @param lease lease object to serialize * @param lease lease object to serialize
* @return an attribute value map representing the lease object * @return an attribute value map representing the lease object
*/ */
public Map<String, AttributeValue> toDynamoRecord(T lease); Map<String, AttributeValue> toDynamoRecord(Lease lease);
/** /**
* Construct a Lease object out of a DynamoDB record. * Construct a Lease object out of a DynamoDB record.
@ -45,72 +43,72 @@ public interface LeaseSerializer<T extends Lease> {
* @param dynamoRecord attribute value map from DynamoDB * @param dynamoRecord attribute value map from DynamoDB
* @return a deserialized lease object representing the attribute value map * @return a deserialized lease object representing the attribute value map
*/ */
public T fromDynamoRecord(Map<String, AttributeValue> dynamoRecord); Lease fromDynamoRecord(Map<String, AttributeValue> dynamoRecord);
/** /**
* @param lease * @param lease
* @return the attribute value map representing a Lease's hash key given a Lease object. * @return the attribute value map representing a Lease's hash key given a Lease object.
*/ */
public Map<String, AttributeValue> getDynamoHashKey(T lease); Map<String, AttributeValue> getDynamoHashKey(Lease lease);
/** /**
* Special getDynamoHashKey implementation used by ILeaseManager.getLease(). * Special getDynamoHashKey implementation used by {@link LeaseRefresher#getLease(String)}.
* *
* @param leaseKey * @param leaseKey
* @return the attribute value map representing a Lease's hash key given a string. * @return the attribute value map representing a Lease's hash key given a string.
*/ */
public Map<String, AttributeValue> getDynamoHashKey(String leaseKey); Map<String, AttributeValue> getDynamoHashKey(String leaseKey);
/** /**
* @param lease * @param lease
* @return the attribute value map asserting that a lease counter is what we expect. * @return the attribute value map asserting that a lease counter is what we expect.
*/ */
public Map<String, ExpectedAttributeValue> getDynamoLeaseCounterExpectation(T lease); Map<String, ExpectedAttributeValue> getDynamoLeaseCounterExpectation(Lease lease);
/** /**
* @param lease * @param lease
* @return the attribute value map asserting that the lease owner is what we expect. * @return the attribute value map asserting that the lease owner is what we expect.
*/ */
public Map<String, ExpectedAttributeValue> getDynamoLeaseOwnerExpectation(T lease); Map<String, ExpectedAttributeValue> getDynamoLeaseOwnerExpectation(Lease lease);
/** /**
* @return the attribute value map asserting that a lease does not exist. * @return the attribute value map asserting that a lease does not exist.
*/ */
public Map<String, ExpectedAttributeValue> getDynamoNonexistantExpectation(); Map<String, ExpectedAttributeValue> getDynamoNonexistantExpectation();
/** /**
* @param lease * @param lease
* @return the attribute value map that increments a lease counter * @return the attribute value map that increments a lease counter
*/ */
public Map<String, AttributeValueUpdate> getDynamoLeaseCounterUpdate(T lease); Map<String, AttributeValueUpdate> getDynamoLeaseCounterUpdate(Lease lease);
/** /**
* @param lease * @param lease
* @param newOwner * @param newOwner
* @return the attribute value map that takes a lease for a new owner * @return the attribute value map that takes a lease for a new owner
*/ */
public Map<String, AttributeValueUpdate> getDynamoTakeLeaseUpdate(T lease, String newOwner); Map<String, AttributeValueUpdate> getDynamoTakeLeaseUpdate(Lease lease, String newOwner);
/** /**
* @param lease * @param lease
* @return the attribute value map that voids a lease * @return the attribute value map that voids a lease
*/ */
public Map<String, AttributeValueUpdate> getDynamoEvictLeaseUpdate(T lease); Map<String, AttributeValueUpdate> getDynamoEvictLeaseUpdate(Lease lease);
/** /**
* @param lease * @param lease
* @return the attribute value map that updates application-specific data for a lease and increments the lease * @return the attribute value map that updates application-specific data for a lease and increments the lease
* counter * counter
*/ */
public Map<String, AttributeValueUpdate> getDynamoUpdateLeaseUpdate(T lease); Map<String, AttributeValueUpdate> getDynamoUpdateLeaseUpdate(Lease lease);
/** /**
* @return the key schema for creating a DynamoDB table to store leases * @return the key schema for creating a DynamoDB table to store leases
*/ */
public Collection<KeySchemaElement> getKeySchema(); Collection<KeySchemaElement> getKeySchema();
/** /**
* @return attribute definitions for creating a DynamoDB table to store leases * @return attribute definitions for creating a DynamoDB table to store leases
*/ */
public Collection<AttributeDefinition> getAttributeDefinitions(); Collection<AttributeDefinition> getAttributeDefinitions();
} }

View file

@ -18,13 +18,12 @@ import java.util.Map;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.Lease;
/** /**
* ILeaseTaker is used by LeaseCoordinator to take new leases, or leases that other workers fail to renew. Each * ILeaseTaker is used by LeaseCoordinator to take new leases, or leases that other workers fail to renew. Each
* LeaseCoordinator instance corresponds to one worker and uses exactly one ILeaseTaker to take leases for that worker. * LeaseCoordinator instance corresponds to one worker and uses exactly one ILeaseTaker to take leases for that worker.
*/ */
public interface LeaseTaker<T extends Lease> { public interface LeaseTaker {
/** /**
* Compute the set of leases available to be taken and attempt to take them. Lease taking rules are: * Compute the set of leases available to be taken and attempt to take them. Lease taking rules are:
@ -39,11 +38,11 @@ public interface LeaseTaker<T extends Lease> {
* @throws DependencyException on unexpected DynamoDB failures * @throws DependencyException on unexpected DynamoDB failures
* @throws InvalidStateException if lease table does not exist * @throws InvalidStateException if lease table does not exist
*/ */
public abstract Map<String, T> takeLeases() throws DependencyException, InvalidStateException; Map<String, Lease> takeLeases() throws DependencyException, InvalidStateException;
/** /**
* @return workerIdentifier for this LeaseTaker * @return workerIdentifier for this LeaseTaker
*/ */
public abstract String getWorkerIdentifier(); String getWorkerIdentifier();
} }

View file

@ -22,7 +22,7 @@ import java.util.List;
/** /**
* *
*/ */
public interface LeaseManagerProxy { public interface ShardDetector {
List<Shard> listShards(); List<Shard> listShards();
} }

View file

@ -34,9 +34,9 @@ import software.amazon.kinesis.lifecycle.TaskType;
@Slf4j @Slf4j
public class ShardSyncTask implements ITask { public class ShardSyncTask implements ITask {
@NonNull @NonNull
private final LeaseManagerProxy leaseManagerProxy; private final ShardDetector shardDetector;
@NonNull @NonNull
private final LeaseManager<KinesisClientLease> leaseManager; private final LeaseRefresher leaseRefresher;
@NonNull @NonNull
private final InitialPositionInStreamExtended initialPosition; private final InitialPositionInStreamExtended initialPosition;
private final boolean cleanupLeasesUponShardCompletion; private final boolean cleanupLeasesUponShardCompletion;
@ -56,7 +56,7 @@ public class ShardSyncTask implements ITask {
Exception exception = null; Exception exception = null;
try { try {
ShardSyncer.checkAndCreateLeasesForNewShards(leaseManagerProxy, leaseManager, initialPosition, ShardSyncer.checkAndCreateLeasesForNewShards(shardDetector, leaseRefresher, initialPosition,
cleanupLeasesUponShardCompletion, ignoreUnexpectedChildShards); cleanupLeasesUponShardCompletion, ignoreUnexpectedChildShards);
if (shardSyncTaskIdleTimeMillis > 0) { if (shardSyncTaskIdleTimeMillis > 0) {
Thread.sleep(shardSyncTaskIdleTimeMillis); Thread.sleep(shardSyncTaskIdleTimeMillis);

View file

@ -39,9 +39,9 @@ import software.amazon.kinesis.metrics.MetricsCollectingTaskDecorator;
@Slf4j @Slf4j
public class ShardSyncTaskManager { public class ShardSyncTaskManager {
@NonNull @NonNull
private final LeaseManagerProxy leaseManagerProxy; private final ShardDetector shardDetector;
@NonNull @NonNull
private final LeaseManager<KinesisClientLease> leaseManager; private final LeaseRefresher leaseRefresher;
@NonNull @NonNull
private final InitialPositionInStreamExtended initialPositionInStream; private final InitialPositionInStreamExtended initialPositionInStream;
private final boolean cleanupLeasesUponShardCompletion; private final boolean cleanupLeasesUponShardCompletion;
@ -76,8 +76,8 @@ public class ShardSyncTaskManager {
currentTask = currentTask =
new MetricsCollectingTaskDecorator( new MetricsCollectingTaskDecorator(
new ShardSyncTask(leaseManagerProxy, new ShardSyncTask(shardDetector,
leaseManager, leaseRefresher,
initialPositionInStream, initialPositionInStream,
cleanupLeasesUponShardCompletion, cleanupLeasesUponShardCompletion,
ignoreUnexpectedChildShards, ignoreUnexpectedChildShards,

View file

@ -57,20 +57,20 @@ public class ShardSyncer {
private ShardSyncer() { private ShardSyncer() {
} }
static synchronized void bootstrapShardLeases(@NonNull final LeaseManagerProxy leaseManagerProxy, static synchronized void bootstrapShardLeases(@NonNull final ShardDetector shardDetector,
@NonNull final LeaseManager<KinesisClientLease> leaseManager, @NonNull final LeaseRefresher leaseRefresher,
@NonNull final InitialPositionInStreamExtended initialPositionInStream, @NonNull final InitialPositionInStreamExtended initialPositionInStream,
final boolean cleanupLeasesOfCompletedShards, final boolean cleanupLeasesOfCompletedShards,
final boolean ignoreUnexpectedChildShards) final boolean ignoreUnexpectedChildShards)
throws DependencyException, InvalidStateException, ProvisionedThroughputException, KinesisClientLibIOException { throws DependencyException, InvalidStateException, ProvisionedThroughputException, KinesisClientLibIOException {
syncShardLeases(leaseManagerProxy, leaseManager, initialPositionInStream, cleanupLeasesOfCompletedShards, syncShardLeases(shardDetector, leaseRefresher, initialPositionInStream, cleanupLeasesOfCompletedShards,
ignoreUnexpectedChildShards); ignoreUnexpectedChildShards);
} }
/** /**
* Check and create leases for any new shards (e.g. following a reshard operation). * Check and create leases for any new shards (e.g. following a reshard operation).
* *
* @param leaseManager * @param leaseRefresher
* @param initialPositionInStream * @param initialPositionInStream
* @param cleanupLeasesOfCompletedShards * @param cleanupLeasesOfCompletedShards
* @param ignoreUnexpectedChildShards * @param ignoreUnexpectedChildShards
@ -79,29 +79,29 @@ public class ShardSyncer {
* @throws ProvisionedThroughputException * @throws ProvisionedThroughputException
* @throws KinesisClientLibIOException * @throws KinesisClientLibIOException
*/ */
public static synchronized void checkAndCreateLeasesForNewShards(@NonNull final LeaseManagerProxy leaseManagerProxy, public static synchronized void checkAndCreateLeasesForNewShards(@NonNull final ShardDetector shardDetector,
@NonNull final LeaseManager<KinesisClientLease> leaseManager, @NonNull final LeaseRefresher leaseRefresher,
@NonNull final InitialPositionInStreamExtended initialPositionInStream, @NonNull final InitialPositionInStreamExtended initialPositionInStream,
final boolean cleanupLeasesOfCompletedShards, final boolean cleanupLeasesOfCompletedShards,
final boolean ignoreUnexpectedChildShards) final boolean ignoreUnexpectedChildShards)
throws DependencyException, InvalidStateException, ProvisionedThroughputException, KinesisClientLibIOException { throws DependencyException, InvalidStateException, ProvisionedThroughputException, KinesisClientLibIOException {
syncShardLeases(leaseManagerProxy, leaseManager, initialPositionInStream, syncShardLeases(shardDetector, leaseRefresher, initialPositionInStream,
cleanupLeasesOfCompletedShards, ignoreUnexpectedChildShards); cleanupLeasesOfCompletedShards, ignoreUnexpectedChildShards);
} }
static synchronized void checkAndCreateLeasesForNewShards(@NonNull final LeaseManagerProxy leaseManagerProxy, static synchronized void checkAndCreateLeasesForNewShards(@NonNull final ShardDetector shardDetector,
@NonNull final LeaseManager<KinesisClientLease> leaseManager, @NonNull final LeaseRefresher leaseRefresher,
@NonNull final InitialPositionInStreamExtended initialPositionInStream, @NonNull final InitialPositionInStreamExtended initialPositionInStream,
final boolean cleanupLeasesOfCompletedShards) final boolean cleanupLeasesOfCompletedShards)
throws DependencyException, InvalidStateException, ProvisionedThroughputException, KinesisClientLibIOException { throws DependencyException, InvalidStateException, ProvisionedThroughputException, KinesisClientLibIOException {
checkAndCreateLeasesForNewShards(leaseManagerProxy, leaseManager, initialPositionInStream, checkAndCreateLeasesForNewShards(shardDetector, leaseRefresher, initialPositionInStream,
cleanupLeasesOfCompletedShards, false); cleanupLeasesOfCompletedShards, false);
} }
/** /**
* Sync leases with Kinesis shards (e.g. at startup, or when we reach end of a shard). * Sync leases with Kinesis shards (e.g. at startup, or when we reach end of a shard).
* *
* @param leaseManager * @param leaseRefresher
* @param initialPosition * @param initialPosition
* @param cleanupLeasesOfCompletedShards * @param cleanupLeasesOfCompletedShards
* @param ignoreUnexpectedChildShards * @param ignoreUnexpectedChildShards
@ -111,13 +111,13 @@ public class ShardSyncer {
* @throws KinesisClientLibIOException * @throws KinesisClientLibIOException
*/ */
// CHECKSTYLE:OFF CyclomaticComplexity // CHECKSTYLE:OFF CyclomaticComplexity
private static synchronized void syncShardLeases(@NonNull final LeaseManagerProxy leaseManagerProxy, private static synchronized void syncShardLeases(@NonNull final ShardDetector shardDetector,
final LeaseManager<KinesisClientLease> leaseManager, final LeaseRefresher leaseRefresher,
final InitialPositionInStreamExtended initialPosition, final InitialPositionInStreamExtended initialPosition,
final boolean cleanupLeasesOfCompletedShards, final boolean cleanupLeasesOfCompletedShards,
final boolean ignoreUnexpectedChildShards) final boolean ignoreUnexpectedChildShards)
throws DependencyException, InvalidStateException, ProvisionedThroughputException, KinesisClientLibIOException { throws DependencyException, InvalidStateException, ProvisionedThroughputException, KinesisClientLibIOException {
List<Shard> shards = getShardList(leaseManagerProxy); List<Shard> shards = getShardList(shardDetector);
log.debug("Num shards: {}", shards.size()); log.debug("Num shards: {}", shards.size());
Map<String, Shard> shardIdToShardMap = constructShardIdToShardMap(shards); Map<String, Shard> shardIdToShardMap = constructShardIdToShardMap(shards);
@ -127,34 +127,34 @@ public class ShardSyncer {
assertAllParentShardsAreClosed(inconsistentShardIds); assertAllParentShardsAreClosed(inconsistentShardIds);
} }
List<KinesisClientLease> currentLeases = leaseManager.listLeases(); List<Lease> currentLeases = leaseRefresher.listLeases();
List<KinesisClientLease> newLeasesToCreate = determineNewLeasesToCreate(shards, currentLeases, initialPosition, List<Lease> newLeasesToCreate = determineNewLeasesToCreate(shards, currentLeases, initialPosition,
inconsistentShardIds); inconsistentShardIds);
log.debug("Num new leases to create: {}", newLeasesToCreate.size()); log.debug("Num new leases to create: {}", newLeasesToCreate.size());
for (KinesisClientLease lease : newLeasesToCreate) { for (Lease lease : newLeasesToCreate) {
long startTimeMillis = System.currentTimeMillis(); long startTimeMillis = System.currentTimeMillis();
boolean success = false; boolean success = false;
try { try {
leaseManager.createLeaseIfNotExists(lease); leaseRefresher.createLeaseIfNotExists(lease);
success = true; success = true;
} finally { } finally {
MetricsHelper.addSuccessAndLatency("CreateLease", startTimeMillis, success, MetricsLevel.DETAILED); MetricsHelper.addSuccessAndLatency("CreateLease", startTimeMillis, success, MetricsLevel.DETAILED);
} }
} }
List<KinesisClientLease> trackedLeases = new ArrayList<>(); List<Lease> trackedLeases = new ArrayList<>();
if (currentLeases != null) { if (currentLeases != null) {
trackedLeases.addAll(currentLeases); trackedLeases.addAll(currentLeases);
} }
trackedLeases.addAll(newLeasesToCreate); trackedLeases.addAll(newLeasesToCreate);
cleanupGarbageLeases(leaseManagerProxy, shards, trackedLeases, leaseManager); cleanupGarbageLeases(shardDetector, shards, trackedLeases, leaseRefresher);
if (cleanupLeasesOfCompletedShards) { if (cleanupLeasesOfCompletedShards) {
cleanupLeasesOfFinishedShards(currentLeases, cleanupLeasesOfFinishedShards(currentLeases,
shardIdToShardMap, shardIdToShardMap,
shardIdToChildShardIdsMap, shardIdToChildShardIdsMap,
trackedLeases, trackedLeases,
leaseManager); leaseRefresher);
} }
} }
// CHECKSTYLE:ON CyclomaticComplexity // CHECKSTYLE:ON CyclomaticComplexity
@ -200,10 +200,10 @@ public class ShardSyncer {
* @param trackedLeaseList * @param trackedLeaseList
* @return * @return
*/ */
static Map<String, KinesisClientLease> constructShardIdToKCLLeaseMap(List<KinesisClientLease> trackedLeaseList) { static Map<String, Lease> constructShardIdToKCLLeaseMap(List<Lease> trackedLeaseList) {
Map<String, KinesisClientLease> trackedLeasesMap = new HashMap<>(); Map<String, Lease> trackedLeasesMap = new HashMap<>();
for (KinesisClientLease lease : trackedLeaseList) { for (Lease lease : trackedLeaseList) {
trackedLeasesMap.put(lease.getLeaseKey(), lease); trackedLeasesMap.put(lease.leaseKey(), lease);
} }
return trackedLeasesMap; return trackedLeasesMap;
} }
@ -313,8 +313,8 @@ public class ShardSyncer {
return shardIdToChildShardIdsMap; return shardIdToChildShardIdsMap;
} }
private static List<Shard> getShardList(@NonNull final LeaseManagerProxy leaseManagerProxy) throws KinesisClientLibIOException { private static List<Shard> getShardList(@NonNull final ShardDetector shardDetector) throws KinesisClientLibIOException {
List<Shard> shards = leaseManagerProxy.listShards(); List<Shard> shards = shardDetector.listShards();
if (shards == null) { if (shards == null) {
throw new KinesisClientLibIOException( throw new KinesisClientLibIOException(
"Stream is not in ACTIVE OR UPDATING state - will retry getting the shard list."); "Stream is not in ACTIVE OR UPDATING state - will retry getting the shard list.");
@ -369,16 +369,16 @@ public class ShardSyncer {
* @param inconsistentShardIds Set of child shard ids having open parents. * @param inconsistentShardIds Set of child shard ids having open parents.
* @return List of new leases to create sorted by starting sequenceNumber of the corresponding shard * @return List of new leases to create sorted by starting sequenceNumber of the corresponding shard
*/ */
static List<KinesisClientLease> determineNewLeasesToCreate(List<Shard> shards, static List<Lease> determineNewLeasesToCreate(List<Shard> shards,
List<KinesisClientLease> currentLeases, List<Lease> currentLeases,
InitialPositionInStreamExtended initialPosition, InitialPositionInStreamExtended initialPosition,
Set<String> inconsistentShardIds) { Set<String> inconsistentShardIds) {
Map<String, KinesisClientLease> shardIdToNewLeaseMap = new HashMap<String, KinesisClientLease>(); Map<String, Lease> shardIdToNewLeaseMap = new HashMap<>();
Map<String, Shard> shardIdToShardMapOfAllKinesisShards = constructShardIdToShardMap(shards); Map<String, Shard> shardIdToShardMapOfAllKinesisShards = constructShardIdToShardMap(shards);
Set<String> shardIdsOfCurrentLeases = new HashSet<String>(); Set<String> shardIdsOfCurrentLeases = new HashSet<>();
for (KinesisClientLease lease : currentLeases) { for (Lease lease : currentLeases) {
shardIdsOfCurrentLeases.add(lease.getLeaseKey()); shardIdsOfCurrentLeases.add(lease.leaseKey());
log.debug("Existing lease: {}", lease); log.debug("Existing lease: {}", lease);
} }
@ -395,7 +395,7 @@ public class ShardSyncer {
log.info("shardId {} is an inconsistent child. Not creating a lease", shardId); log.info("shardId {} is an inconsistent child. Not creating a lease", shardId);
} else { } else {
log.debug("Need to create a lease for shardId {}", shardId); log.debug("Need to create a lease for shardId {}", shardId);
KinesisClientLease newLease = newKCLLease(shard); Lease newLease = newKCLLease(shard);
boolean isDescendant = boolean isDescendant =
checkIfDescendantAndAddNewLeasesForAncestors(shardId, checkIfDescendantAndAddNewLeasesForAncestors(shardId,
initialPosition, initialPosition,
@ -429,18 +429,18 @@ public class ShardSyncer {
*/ */
if (isDescendant && !initialPosition.getInitialPositionInStream() if (isDescendant && !initialPosition.getInitialPositionInStream()
.equals(InitialPositionInStream.AT_TIMESTAMP)) { .equals(InitialPositionInStream.AT_TIMESTAMP)) {
newLease.setCheckpoint(ExtendedSequenceNumber.TRIM_HORIZON); newLease.checkpoint(ExtendedSequenceNumber.TRIM_HORIZON);
} else { } else {
newLease.setCheckpoint(convertToCheckpoint(initialPosition)); newLease.checkpoint(convertToCheckpoint(initialPosition));
} }
log.debug("Set checkpoint of {} to {}", newLease.getLeaseKey(), newLease.getCheckpoint()); log.debug("Set checkpoint of {} to {}", newLease.leaseKey(), newLease.checkpoint());
shardIdToNewLeaseMap.put(shardId, newLease); shardIdToNewLeaseMap.put(shardId, newLease);
} }
} }
List<KinesisClientLease> newLeasesToCreate = new ArrayList<KinesisClientLease>(); List<Lease> newLeasesToCreate = new ArrayList<>();
newLeasesToCreate.addAll(shardIdToNewLeaseMap.values()); newLeasesToCreate.addAll(shardIdToNewLeaseMap.values());
Comparator<? super KinesisClientLease> startingSequenceNumberComparator = Comparator<Lease> startingSequenceNumberComparator =
new StartingSequenceNumberAndShardIdBasedComparator(shardIdToShardMapOfAllKinesisShards); new StartingSequenceNumberAndShardIdBasedComparator(shardIdToShardMapOfAllKinesisShards);
Collections.sort(newLeasesToCreate, startingSequenceNumberComparator); Collections.sort(newLeasesToCreate, startingSequenceNumberComparator);
return newLeasesToCreate; return newLeasesToCreate;
@ -450,10 +450,10 @@ public class ShardSyncer {
* Determine new leases to create and their initial checkpoint. * Determine new leases to create and their initial checkpoint.
* Note: Package level access only for testing purposes. * Note: Package level access only for testing purposes.
*/ */
static List<KinesisClientLease> determineNewLeasesToCreate(List<Shard> shards, static List<Lease> determineNewLeasesToCreate(List<Shard> shards,
List<KinesisClientLease> currentLeases, List<Lease> currentLeases,
InitialPositionInStreamExtended initialPosition) { InitialPositionInStreamExtended initialPosition) {
Set<String> inconsistentShardIds = new HashSet<String>(); Set<String> inconsistentShardIds = new HashSet<>();
return determineNewLeasesToCreate(shards, currentLeases, initialPosition, inconsistentShardIds); return determineNewLeasesToCreate(shards, currentLeases, initialPosition, inconsistentShardIds);
} }
@ -477,7 +477,7 @@ public class ShardSyncer {
InitialPositionInStreamExtended initialPosition, InitialPositionInStreamExtended initialPosition,
Set<String> shardIdsOfCurrentLeases, Set<String> shardIdsOfCurrentLeases,
Map<String, Shard> shardIdToShardMapOfAllKinesisShards, Map<String, Shard> shardIdToShardMapOfAllKinesisShards,
Map<String, KinesisClientLease> shardIdToLeaseMapOfNewShards, Map<String, Lease> shardIdToLeaseMapOfNewShards,
Map<String, Boolean> memoizationContext) { Map<String, Boolean> memoizationContext) {
Boolean previousValue = memoizationContext.get(shardId); Boolean previousValue = memoizationContext.get(shardId);
@ -520,7 +520,7 @@ public class ShardSyncer {
for (String parentShardId : parentShardIds) { for (String parentShardId : parentShardIds) {
if (!shardIdsOfCurrentLeases.contains(parentShardId)) { if (!shardIdsOfCurrentLeases.contains(parentShardId)) {
log.debug("Need to create a lease for shardId {}", parentShardId); log.debug("Need to create a lease for shardId {}", parentShardId);
KinesisClientLease lease = shardIdToLeaseMapOfNewShards.get(parentShardId); Lease lease = shardIdToLeaseMapOfNewShards.get(parentShardId);
if (lease == null) { if (lease == null) {
lease = newKCLLease(shardIdToShardMapOfAllKinesisShards.get(parentShardId)); lease = newKCLLease(shardIdToShardMapOfAllKinesisShards.get(parentShardId));
shardIdToLeaseMapOfNewShards.put(parentShardId, lease); shardIdToLeaseMapOfNewShards.put(parentShardId, lease);
@ -529,9 +529,9 @@ public class ShardSyncer {
if (descendantParentShardIds.contains(parentShardId) if (descendantParentShardIds.contains(parentShardId)
&& !initialPosition.getInitialPositionInStream() && !initialPosition.getInitialPositionInStream()
.equals(InitialPositionInStream.AT_TIMESTAMP)) { .equals(InitialPositionInStream.AT_TIMESTAMP)) {
lease.setCheckpoint(ExtendedSequenceNumber.TRIM_HORIZON); lease.checkpoint(ExtendedSequenceNumber.TRIM_HORIZON);
} else { } else {
lease.setCheckpoint(convertToCheckpoint(initialPosition)); lease.checkpoint(convertToCheckpoint(initialPosition));
} }
} }
} }
@ -584,16 +584,16 @@ public class ShardSyncer {
* * the parentShardIds listed in the lease are also not present in the list of Kinesis shards. * * the parentShardIds listed in the lease are also not present in the list of Kinesis shards.
* @param shards List of all Kinesis shards (assumed to be a consistent snapshot - when stream is in Active state). * @param shards List of all Kinesis shards (assumed to be a consistent snapshot - when stream is in Active state).
* @param trackedLeases List of * @param trackedLeases List of
* @param leaseManager * @param leaseRefresher
* @throws KinesisClientLibIOException Thrown if we couldn't get a fresh shard list from Kinesis. * @throws KinesisClientLibIOException Thrown if we couldn't get a fresh shard list from Kinesis.
* @throws ProvisionedThroughputException * @throws ProvisionedThroughputException
* @throws InvalidStateException * @throws InvalidStateException
* @throws DependencyException * @throws DependencyException
*/ */
private static void cleanupGarbageLeases(@NonNull final LeaseManagerProxy leaseManagerProxy, private static void cleanupGarbageLeases(@NonNull final ShardDetector shardDetector,
final List<Shard> shards, final List<Shard> shards,
final List<KinesisClientLease> trackedLeases, final List<Lease> trackedLeases,
final LeaseManager<KinesisClientLease> leaseManager) final LeaseRefresher leaseRefresher)
throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException { throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException {
Set<String> kinesisShards = new HashSet<>(); Set<String> kinesisShards = new HashSet<>();
for (Shard shard : shards) { for (Shard shard : shards) {
@ -601,8 +601,8 @@ public class ShardSyncer {
} }
// Check if there are leases for non-existent shards // Check if there are leases for non-existent shards
List<KinesisClientLease> garbageLeases = new ArrayList<>(); List<Lease> garbageLeases = new ArrayList<>();
for (KinesisClientLease lease : trackedLeases) { for (Lease lease : trackedLeases) {
if (isCandidateForCleanup(lease, kinesisShards)) { if (isCandidateForCleanup(lease, kinesisShards)) {
garbageLeases.add(lease); garbageLeases.add(lease);
} }
@ -611,16 +611,16 @@ public class ShardSyncer {
if (!garbageLeases.isEmpty()) { if (!garbageLeases.isEmpty()) {
log.info("Found {} candidate leases for cleanup. Refreshing list of" log.info("Found {} candidate leases for cleanup. Refreshing list of"
+ " Kinesis shards to pick up recent/latest shards", garbageLeases.size()); + " Kinesis shards to pick up recent/latest shards", garbageLeases.size());
List<Shard> currentShardList = getShardList(leaseManagerProxy); List<Shard> currentShardList = getShardList(shardDetector);
Set<String> currentKinesisShardIds = new HashSet<>(); Set<String> currentKinesisShardIds = new HashSet<>();
for (Shard shard : currentShardList) { for (Shard shard : currentShardList) {
currentKinesisShardIds.add(shard.getShardId()); currentKinesisShardIds.add(shard.getShardId());
} }
for (KinesisClientLease lease : garbageLeases) { for (Lease lease : garbageLeases) {
if (isCandidateForCleanup(lease, currentKinesisShardIds)) { if (isCandidateForCleanup(lease, currentKinesisShardIds)) {
log.info("Deleting lease for shard {} as it is not present in Kinesis stream.", lease.getLeaseKey()); log.info("Deleting lease for shard {} as it is not present in Kinesis stream.", lease.leaseKey());
leaseManager.deleteLease(lease); leaseRefresher.deleteLease(lease);
} }
} }
} }
@ -637,15 +637,15 @@ public class ShardSyncer {
* @throws KinesisClientLibIOException Thrown if currentKinesisShardIds contains a parent shard but not the child * @throws KinesisClientLibIOException Thrown if currentKinesisShardIds contains a parent shard but not the child
* shard (we are evaluating for deletion). * shard (we are evaluating for deletion).
*/ */
static boolean isCandidateForCleanup(KinesisClientLease lease, Set<String> currentKinesisShardIds) static boolean isCandidateForCleanup(Lease lease, Set<String> currentKinesisShardIds)
throws KinesisClientLibIOException { throws KinesisClientLibIOException {
boolean isCandidateForCleanup = true; boolean isCandidateForCleanup = true;
if (currentKinesisShardIds.contains(lease.getLeaseKey())) { if (currentKinesisShardIds.contains(lease.leaseKey())) {
isCandidateForCleanup = false; isCandidateForCleanup = false;
} else { } else {
log.info("Found lease for non-existent shard: {}. Checking its parent shards", lease.getLeaseKey()); log.info("Found lease for non-existent shard: {}. Checking its parent shards", lease.leaseKey());
Set<String> parentShardIds = lease.getParentShardIds(); Set<String> parentShardIds = lease.parentShardIds();
for (String parentShardId : parentShardIds) { for (String parentShardId : parentShardIds) {
// Throw an exception if the parent shard exists (but the child does not). // Throw an exception if the parent shard exists (but the child does not).
@ -653,7 +653,7 @@ public class ShardSyncer {
if (currentKinesisShardIds.contains(parentShardId)) { if (currentKinesisShardIds.contains(parentShardId)) {
String message = String message =
"Parent shard " + parentShardId + " exists but not the child shard " "Parent shard " + parentShardId + " exists but not the child shard "
+ lease.getLeaseKey(); + lease.leaseKey();
log.info(message); log.info(message);
throw new KinesisClientLibIOException(message); throw new KinesisClientLibIOException(message);
} }
@ -674,23 +674,23 @@ public class ShardSyncer {
* @param shardIdToShardMap Map of shardId->Shard (assumed to include all Kinesis shards) * @param shardIdToShardMap Map of shardId->Shard (assumed to include all Kinesis shards)
* @param shardIdToChildShardIdsMap Map of shardId->childShardIds (assumed to include all Kinesis shards) * @param shardIdToChildShardIdsMap Map of shardId->childShardIds (assumed to include all Kinesis shards)
* @param trackedLeases List of all leases we are tracking. * @param trackedLeases List of all leases we are tracking.
* @param leaseManager Lease manager (will be used to delete leases) * @param leaseRefresher Lease refresher (will be used to delete leases)
* @throws DependencyException * @throws DependencyException
* @throws InvalidStateException * @throws InvalidStateException
* @throws ProvisionedThroughputException * @throws ProvisionedThroughputException
* @throws KinesisClientLibIOException * @throws KinesisClientLibIOException
*/ */
private static synchronized void cleanupLeasesOfFinishedShards(Collection<KinesisClientLease> currentLeases, private static synchronized void cleanupLeasesOfFinishedShards(Collection<Lease> currentLeases,
Map<String, Shard> shardIdToShardMap, Map<String, Shard> shardIdToShardMap,
Map<String, Set<String>> shardIdToChildShardIdsMap, Map<String, Set<String>> shardIdToChildShardIdsMap,
List<KinesisClientLease> trackedLeases, List<Lease> trackedLeases,
LeaseManager<KinesisClientLease> leaseManager) LeaseRefresher leaseRefresher)
throws DependencyException, InvalidStateException, ProvisionedThroughputException, KinesisClientLibIOException { throws DependencyException, InvalidStateException, ProvisionedThroughputException, KinesisClientLibIOException {
Set<String> shardIdsOfClosedShards = new HashSet<>(); Set<String> shardIdsOfClosedShards = new HashSet<>();
List<KinesisClientLease> leasesOfClosedShards = new ArrayList<>(); List<Lease> leasesOfClosedShards = new ArrayList<>();
for (KinesisClientLease lease : currentLeases) { for (Lease lease : currentLeases) {
if (lease.getCheckpoint().equals(ExtendedSequenceNumber.SHARD_END)) { if (lease.checkpoint().equals(ExtendedSequenceNumber.SHARD_END)) {
shardIdsOfClosedShards.add(lease.getLeaseKey()); shardIdsOfClosedShards.add(lease.leaseKey());
leasesOfClosedShards.add(lease); leasesOfClosedShards.add(lease);
} }
} }
@ -699,16 +699,16 @@ public class ShardSyncer {
assertClosedShardsAreCoveredOrAbsent(shardIdToShardMap, assertClosedShardsAreCoveredOrAbsent(shardIdToShardMap,
shardIdToChildShardIdsMap, shardIdToChildShardIdsMap,
shardIdsOfClosedShards); shardIdsOfClosedShards);
Comparator<? super KinesisClientLease> startingSequenceNumberComparator = Comparator<? super Lease> startingSequenceNumberComparator =
new StartingSequenceNumberAndShardIdBasedComparator(shardIdToShardMap); new StartingSequenceNumberAndShardIdBasedComparator(shardIdToShardMap);
Collections.sort(leasesOfClosedShards, startingSequenceNumberComparator); Collections.sort(leasesOfClosedShards, startingSequenceNumberComparator);
Map<String, KinesisClientLease> trackedLeaseMap = constructShardIdToKCLLeaseMap(trackedLeases); Map<String, Lease> trackedLeaseMap = constructShardIdToKCLLeaseMap(trackedLeases);
for (KinesisClientLease leaseOfClosedShard : leasesOfClosedShards) { for (Lease leaseOfClosedShard : leasesOfClosedShards) {
String closedShardId = leaseOfClosedShard.getLeaseKey(); String closedShardId = leaseOfClosedShard.leaseKey();
Set<String> childShardIds = shardIdToChildShardIdsMap.get(closedShardId); Set<String> childShardIds = shardIdToChildShardIdsMap.get(closedShardId);
if ((closedShardId != null) && (childShardIds != null) && (!childShardIds.isEmpty())) { if ((closedShardId != null) && (childShardIds != null) && (!childShardIds.isEmpty())) {
cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, leaseManager); cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, leaseRefresher);
} }
} }
} }
@ -722,33 +722,33 @@ public class ShardSyncer {
* *
* @param closedShardId Identifies the closed shard * @param closedShardId Identifies the closed shard
* @param childShardIds ShardIds of children of the closed shard * @param childShardIds ShardIds of children of the closed shard
* @param trackedLeases shardId->KinesisClientLease map with all leases we are tracking (should not be null) * @param trackedLeases shardId->Lease map with all leases we are tracking (should not be null)
* @param leaseManager * @param leaseRefresher
* @throws ProvisionedThroughputException * @throws ProvisionedThroughputException
* @throws InvalidStateException * @throws InvalidStateException
* @throws DependencyException * @throws DependencyException
*/ */
static synchronized void cleanupLeaseForClosedShard(String closedShardId, static synchronized void cleanupLeaseForClosedShard(String closedShardId,
Set<String> childShardIds, Set<String> childShardIds,
Map<String, KinesisClientLease> trackedLeases, Map<String, Lease> trackedLeases,
LeaseManager<KinesisClientLease> leaseManager) LeaseRefresher leaseRefresher)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
KinesisClientLease leaseForClosedShard = trackedLeases.get(closedShardId); Lease leaseForClosedShard = trackedLeases.get(closedShardId);
List<KinesisClientLease> childShardLeases = new ArrayList<>(); List<Lease> childShardLeases = new ArrayList<>();
for (String childShardId : childShardIds) { for (String childShardId : childShardIds) {
KinesisClientLease childLease = trackedLeases.get(childShardId); Lease childLease = trackedLeases.get(childShardId);
if (childLease != null) { if (childLease != null) {
childShardLeases.add(childLease); childShardLeases.add(childLease);
} }
} }
if ((leaseForClosedShard != null) if ((leaseForClosedShard != null)
&& (leaseForClosedShard.getCheckpoint().equals(ExtendedSequenceNumber.SHARD_END)) && (leaseForClosedShard.checkpoint().equals(ExtendedSequenceNumber.SHARD_END))
&& (childShardLeases.size() == childShardIds.size())) { && (childShardLeases.size() == childShardIds.size())) {
boolean okayToDelete = true; boolean okayToDelete = true;
for (KinesisClientLease lease : childShardLeases) { for (Lease lease : childShardLeases) {
if (lease.getCheckpoint().equals(ExtendedSequenceNumber.TRIM_HORIZON)) { if (lease.checkpoint().equals(ExtendedSequenceNumber.TRIM_HORIZON)) {
okayToDelete = false; okayToDelete = false;
break; break;
} }
@ -756,22 +756,22 @@ public class ShardSyncer {
if (okayToDelete) { if (okayToDelete) {
log.info("Deleting lease for shard {} as it has been completely processed and processing of child " log.info("Deleting lease for shard {} as it has been completely processed and processing of child "
+ "shards has begun.", leaseForClosedShard.getLeaseKey()); + "shards has begun.", leaseForClosedShard.leaseKey());
leaseManager.deleteLease(leaseForClosedShard); leaseRefresher.deleteLease(leaseForClosedShard);
} }
} }
} }
/** /**
* Helper method to create a new KinesisClientLease POJO for a shard. * Helper method to create a new Lease POJO for a shard.
* Note: Package level access only for testing purposes * Note: Package level access only for testing purposes
* *
* @param shard * @param shard
* @return * @return
*/ */
public static KinesisClientLease newKCLLease(Shard shard) { public static Lease newKCLLease(Shard shard) {
KinesisClientLease newLease = new KinesisClientLease(); Lease newLease = new Lease();
newLease.setLeaseKey(shard.getShardId()); newLease.leaseKey(shard.getShardId());
List<String> parentShardIds = new ArrayList<String>(2); List<String> parentShardIds = new ArrayList<String>(2);
if (shard.getParentShardId() != null) { if (shard.getParentShardId() != null) {
parentShardIds.add(shard.getParentShardId()); parentShardIds.add(shard.getParentShardId());
@ -779,8 +779,8 @@ public class ShardSyncer {
if (shard.getAdjacentParentShardId() != null) { if (shard.getAdjacentParentShardId() != null) {
parentShardIds.add(shard.getAdjacentParentShardId()); parentShardIds.add(shard.getAdjacentParentShardId());
} }
newLease.setParentShardIds(parentShardIds); newLease.parentShardIds(parentShardIds);
newLease.setOwnerSwitchesSinceCheckpoint(0L); newLease.ownerSwitchesSinceCheckpoint(0L);
return newLease; return newLease;
} }
@ -835,8 +835,7 @@ public class ShardSyncer {
/** Helper class to compare leases based on starting sequence number of the corresponding shards. /** Helper class to compare leases based on starting sequence number of the corresponding shards.
* *
*/ */
private static class StartingSequenceNumberAndShardIdBasedComparator implements Comparator<KinesisClientLease>, private static class StartingSequenceNumberAndShardIdBasedComparator implements Comparator<Lease>, Serializable {
Serializable {
private static final long serialVersionUID = 1L; private static final long serialVersionUID = 1L;
@ -859,10 +858,10 @@ public class ShardSyncer {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public int compare(KinesisClientLease lease1, KinesisClientLease lease2) { public int compare(Lease lease1, Lease lease2) {
int result = 0; int result = 0;
String shardId1 = lease1.getLeaseKey(); String shardId1 = lease1.leaseKey();
String shardId2 = lease2.getLeaseKey(); String shardId2 = lease2.leaseKey();
Shard shard1 = shardIdToShardMap.get(shardId1); Shard shard1 = shardIdToShardMap.get(shardId1);
Shard shard2 = shardIdToShardMap.get(shardId2); Shard shard2 = shardIdToShardMap.get(shardId2);

View file

@ -0,0 +1,357 @@
/*
* Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Amazon Software License (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/asl/
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*/
package software.amazon.kinesis.leases.dynamodb;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedTransferQueue;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import com.amazonaws.util.CollectionUtils;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.LeaseCoordinator;
import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.LeaseRenewer;
import software.amazon.kinesis.leases.LeaseTaker;
import software.amazon.kinesis.leases.ShardInfo;
import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.LeasingException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
import software.amazon.kinesis.metrics.IMetricsFactory;
import software.amazon.kinesis.metrics.IMetricsScope;
import software.amazon.kinesis.metrics.MetricsHelper;
import software.amazon.kinesis.metrics.MetricsLevel;
/**
* LeaseCoordinator abstracts away LeaseTaker and LeaseRenewer from the application code that's using leasing. It owns
* the scheduling of the two previously mentioned components as well as informing LeaseRenewer when LeaseTaker takes new
* leases.
*
*/
@Slf4j
public class DynamoDBLeaseCoordinator implements LeaseCoordinator {
/*
* Name of the dimension used when setting worker identifier on IMetricsScopes. Exposed so that users of this class
* can easily create InterceptingMetricsFactories that rename this dimension to suit the destination metrics system.
*/
private static final String WORKER_IDENTIFIER_METRIC = "WorkerIdentifier";
// Time to wait for in-flight Runnables to finish when calling .stop();
private static final long STOP_WAIT_TIME_MILLIS = 2000L;
private static final long DEFAULT_INITIAL_LEASE_TABLE_READ_CAPACITY = 10L;
private static final long DEFAULT_INITIAL_LEASE_TABLE_WRITE_CAPACITY = 10L;
private static final ThreadFactory LEASE_COORDINATOR_THREAD_FACTORY = new ThreadFactoryBuilder()
.setNameFormat("LeaseCoordinator-%04d").setDaemon(true).build();
private static final ThreadFactory LEASE_RENEWAL_THREAD_FACTORY = new ThreadFactoryBuilder()
.setNameFormat("LeaseRenewer-%04d").setDaemon(true).build();
private final LeaseRenewer leaseRenewer;
private final LeaseTaker leaseTaker;
private final long renewerIntervalMillis;
private final long takerIntervalMillis;
private final ExecutorService leaseRenewalThreadpool;
private final LeaseRefresher leaseRefresher;
private final Object shutdownLock = new Object();
protected final IMetricsFactory metricsFactory;
private long initialLeaseTableReadCapacity = DEFAULT_INITIAL_LEASE_TABLE_READ_CAPACITY;
private long initialLeaseTableWriteCapacity = DEFAULT_INITIAL_LEASE_TABLE_WRITE_CAPACITY;
private ScheduledExecutorService leaseCoordinatorThreadPool;
private ScheduledFuture<?> takerFuture;
private volatile boolean running = false;
/**
* Constructor.
*
* @param leaseRefresher LeaseRefresher instance to use
* @param workerIdentifier Identifies the worker (e.g. useful to track lease ownership)
* @param leaseDurationMillis Duration of a lease
* @param epsilonMillis Allow for some variance when calculating lease expirations
* @param maxLeasesForWorker Max leases this Worker can handle at a time
* @param maxLeasesToStealAtOneTime Steal up to these many leases at a time (for load balancing)
* @param metricsFactory Used to publish metrics about lease operations
*/
public DynamoDBLeaseCoordinator(final LeaseRefresher leaseRefresher,
final String workerIdentifier,
final long leaseDurationMillis,
final long epsilonMillis,
final int maxLeasesForWorker,
final int maxLeasesToStealAtOneTime,
final int maxLeaseRenewerThreadCount,
final IMetricsFactory metricsFactory) {
this.leaseRefresher = leaseRefresher;
this.leaseRenewalThreadpool = getLeaseRenewalExecutorService(maxLeaseRenewerThreadCount);
this.leaseTaker = new DynamoDBLeaseTaker(leaseRefresher, workerIdentifier, leaseDurationMillis)
.withMaxLeasesForWorker(maxLeasesForWorker)
.withMaxLeasesToStealAtOneTime(maxLeasesToStealAtOneTime);
this.leaseRenewer = new DynamoDBLeaseRenewer(
leaseRefresher, workerIdentifier, leaseDurationMillis, leaseRenewalThreadpool);
this.renewerIntervalMillis = leaseDurationMillis / 3 - epsilonMillis;
this.takerIntervalMillis = (leaseDurationMillis + epsilonMillis) * 2;
this.metricsFactory = metricsFactory;
log.info("With failover time {} ms and epsilon {} ms, LeaseCoordinator will renew leases every {} ms, take"
+ "leases every {} ms, process maximum of {} leases and steal {} lease(s) at a time.",
leaseDurationMillis,
epsilonMillis,
renewerIntervalMillis,
takerIntervalMillis,
maxLeasesForWorker,
maxLeasesToStealAtOneTime);
}
private class TakerRunnable implements Runnable {
@Override
public void run() {
try {
runLeaseTaker();
} catch (LeasingException e) {
log.error("LeasingException encountered in lease taking thread", e);
} catch (Throwable t) {
log.error("Throwable encountered in lease taking thread", t);
}
}
}
private class RenewerRunnable implements Runnable {
@Override
public void run() {
try {
runLeaseRenewer();
} catch (LeasingException e) {
log.error("LeasingException encountered in lease renewing thread", e);
} catch (Throwable t) {
log.error("Throwable encountered in lease renewing thread", t);
}
}
}
@Override
public void initialize() throws ProvisionedThroughputException, DependencyException, IllegalStateException {
final boolean newTableCreated =
leaseRefresher.createLeaseTableIfNotExists(initialLeaseTableReadCapacity, initialLeaseTableWriteCapacity);
if (newTableCreated) {
log.info("Created new lease table for coordinator with initial read capacity of {} and write capacity of {}.",
initialLeaseTableReadCapacity, initialLeaseTableWriteCapacity);
}
// Need to wait for table in active state.
final long secondsBetweenPolls = 10L;
final long timeoutSeconds = 600L;
final boolean isTableActive = leaseRefresher.waitUntilLeaseTableExists(secondsBetweenPolls, timeoutSeconds);
if (!isTableActive) {
throw new DependencyException(new IllegalStateException("Creating table timeout"));
}
}
@Override
public void start() throws DependencyException, InvalidStateException, ProvisionedThroughputException {
leaseRenewer.initialize();
// 2 because we know we'll have at most 2 concurrent tasks at a time.
leaseCoordinatorThreadPool = Executors.newScheduledThreadPool(2, LEASE_COORDINATOR_THREAD_FACTORY);
// Taker runs with fixed DELAY because we want it to run slower in the event of performance degredation.
takerFuture = leaseCoordinatorThreadPool.scheduleWithFixedDelay(new TakerRunnable(),
0L,
takerIntervalMillis,
TimeUnit.MILLISECONDS);
// Renewer runs at fixed INTERVAL because we want it to run at the same rate in the event of degredation.
leaseCoordinatorThreadPool.scheduleAtFixedRate(new RenewerRunnable(),
0L,
renewerIntervalMillis,
TimeUnit.MILLISECONDS);
running = true;
}
@Override
public void runLeaseTaker() throws DependencyException, InvalidStateException {
IMetricsScope scope = MetricsHelper.startScope(metricsFactory, "TakeLeases");
long startTime = System.currentTimeMillis();
boolean success = false;
try {
Map<String, Lease> takenLeases = leaseTaker.takeLeases();
// Only add taken leases to renewer if coordinator is still running.
synchronized (shutdownLock) {
if (running) {
leaseRenewer.addLeasesToRenew(takenLeases.values());
}
}
success = true;
} finally {
scope.addDimension(WORKER_IDENTIFIER_METRIC, workerIdentifier());
MetricsHelper.addSuccessAndLatency(startTime, success, MetricsLevel.SUMMARY);
MetricsHelper.endScope();
}
}
@Override
public void runLeaseRenewer() throws DependencyException, InvalidStateException {
IMetricsScope scope = MetricsHelper.startScope(metricsFactory, "RenewAllLeases");
long startTime = System.currentTimeMillis();
boolean success = false;
try {
leaseRenewer.renewLeases();
success = true;
} finally {
scope.addDimension(WORKER_IDENTIFIER_METRIC, workerIdentifier());
MetricsHelper.addSuccessAndLatency(startTime, success, MetricsLevel.SUMMARY);
MetricsHelper.endScope();
}
}
@Override
public Collection<Lease> getAssignments() {
return leaseRenewer.getCurrentlyHeldLeases().values();
}
@Override
public Lease getCurrentlyHeldLease(String leaseKey) {
return leaseRenewer.getCurrentlyHeldLease(leaseKey);
}
@Override
public String workerIdentifier() {
return leaseTaker.getWorkerIdentifier();
}
@Override
public LeaseRefresher leaseRefresher() {
return leaseRefresher;
}
@Override
public void stop() {
if (leaseCoordinatorThreadPool != null) {
leaseCoordinatorThreadPool.shutdown();
try {
if (leaseCoordinatorThreadPool.awaitTermination(STOP_WAIT_TIME_MILLIS, TimeUnit.MILLISECONDS)) {
log.info("Worker {} has successfully stopped lease-tracking threads",
leaseTaker.getWorkerIdentifier());
} else {
leaseCoordinatorThreadPool.shutdownNow();
log.info("Worker {} stopped lease-tracking threads {} ms after stop",
leaseTaker.getWorkerIdentifier(),
STOP_WAIT_TIME_MILLIS);
}
} catch (InterruptedException e) {
log.debug("Encountered InterruptedException when awaiting threadpool termination");
}
} else {
log.debug("Threadpool was null, no need to shutdown/terminate threadpool.");
}
leaseRenewalThreadpool.shutdownNow();
synchronized (shutdownLock) {
leaseRenewer.clearCurrentlyHeldLeases();
running = false;
}
}
@Override
public void stopLeaseTaker() {
takerFuture.cancel(false);
}
@Override
public void dropLease(final Lease lease) {
synchronized (shutdownLock) {
if (lease != null) {
leaseRenewer.dropLease(lease);
}
}
}
@Override
public boolean isRunning() {
return running;
}
@Override
public boolean updateLease(final Lease lease, final UUID concurrencyToken)
throws DependencyException, InvalidStateException, ProvisionedThroughputException {
return leaseRenewer.updateLease(lease, concurrencyToken);
}
/**
* Returns executor service that should be used for lease renewal.
* @param maximumPoolSize Maximum allowed thread pool size
* @return Executor service that should be used for lease renewal.
*/
private static ExecutorService getLeaseRenewalExecutorService(int maximumPoolSize) {
int coreLeaseCount = Math.max(maximumPoolSize / 4, 2);
return new ThreadPoolExecutor(coreLeaseCount, maximumPoolSize, 60, TimeUnit.SECONDS,
new LinkedTransferQueue<>(), LEASE_RENEWAL_THREAD_FACTORY);
}
@Override
public List<ShardInfo> getCurrentAssignments() {
Collection<Lease> leases = getAssignments();
return convertLeasesToAssignments(leases);
}
private static List<ShardInfo> convertLeasesToAssignments(final Collection<Lease> leases) {
if (CollectionUtils.isNullOrEmpty(leases)) {
return Collections.emptyList();
}
return leases.stream().map(DynamoDBLeaseCoordinator::convertLeaseToAssignment).collect(Collectors.toList());
}
public static ShardInfo convertLeaseToAssignment(final Lease lease) {
return new ShardInfo(lease.leaseKey(), lease.concurrencyToken().toString(), lease.parentShardIds(),
lease.checkpoint());
}
@Override
public DynamoDBLeaseCoordinator initialLeaseTableReadCapacity(long readCapacity) {
if (readCapacity <= 0) {
throw new IllegalArgumentException("readCapacity should be >= 1");
}
this.initialLeaseTableReadCapacity = readCapacity;
return this;
}
@Override
public DynamoDBLeaseCoordinator initialLeaseTableWriteCapacity(long writeCapacity) {
if (writeCapacity <= 0) {
throw new IllegalArgumentException("writeCapacity should be >= 1");
}
this.initialLeaseTableWriteCapacity = writeCapacity;
return this;
}
}

View file

@ -13,7 +13,7 @@
* permissions and limitations under the License. * permissions and limitations under the License.
*/ */
package software.amazon.kinesis.leases; package software.amazon.kinesis.leases.dynamodb;
import java.util.concurrent.ExecutorService; import java.util.concurrent.ExecutorService;
@ -23,6 +23,11 @@ import com.amazonaws.services.kinesis.clientlibrary.lib.worker.InitialPositionIn
import lombok.Data; import lombok.Data;
import lombok.NonNull; import lombok.NonNull;
import software.amazon.kinesis.leases.KinesisShardDetector;
import software.amazon.kinesis.leases.LeaseCoordinator;
import software.amazon.kinesis.leases.LeaseManagementFactory;
import software.amazon.kinesis.leases.ShardDetector;
import software.amazon.kinesis.leases.ShardSyncTaskManager;
import software.amazon.kinesis.metrics.IMetricsFactory; import software.amazon.kinesis.metrics.IMetricsFactory;
/** /**
@ -60,29 +65,7 @@ public class DynamoDBLeaseManagementFactory implements LeaseManagementFactory {
@Override @Override
public LeaseCoordinator createLeaseCoordinator() { public LeaseCoordinator createLeaseCoordinator() {
return createKinesisClientLibLeaseCoordinator(); return new DynamoDBLeaseCoordinator(this.createLeaseRefresher(),
}
@Override
public ShardSyncTaskManager createShardSyncTaskManager() {
return new ShardSyncTaskManager(this.createLeaseManagerProxy(),
this.createLeaseManager(),
initialPositionInStream,
cleanupLeasesUponShardCompletion,
ignoreUnexpectedChildShards,
shardSyncIntervalMillis,
executorService,
metricsFactory);
}
@Override
public DynamoDBLeaseManager<KinesisClientLease> createLeaseManager() {
return new KinesisClientDynamoDBLeaseManager(tableName, amazonDynamoDB, consistentReads);
}
@Override
public KinesisClientLibLeaseCoordinator createKinesisClientLibLeaseCoordinator() {
return new KinesisClientLibLeaseCoordinator(this.createLeaseManager(),
workerIdentifier, workerIdentifier,
failoverTimeMillis, failoverTimeMillis,
epsilonMillis, epsilonMillis,
@ -93,8 +76,25 @@ public class DynamoDBLeaseManagementFactory implements LeaseManagementFactory {
} }
@Override @Override
public LeaseManagerProxy createLeaseManagerProxy() { public ShardSyncTaskManager createShardSyncTaskManager() {
return new KinesisLeaseManagerProxy(amazonKinesis, streamName, listShardsBackoffTimeMillis, return new ShardSyncTaskManager(this.createShardDetector(),
this.createLeaseRefresher(),
initialPositionInStream,
cleanupLeasesUponShardCompletion,
ignoreUnexpectedChildShards,
shardSyncIntervalMillis,
executorService,
metricsFactory);
}
@Override
public DynamoDBLeaseRefresher createLeaseRefresher() {
return new DynamoDBLeaseRefresher(tableName, amazonDynamoDB, new DynamoDBLeaseSerializer(), consistentReads);
}
@Override
public ShardDetector createShardDetector() {
return new KinesisShardDetector(amazonKinesis, streamName, listShardsBackoffTimeMillis,
maxListShardsRetryAttempts); maxListShardsRetryAttempts);
} }
} }

View file

@ -12,7 +12,7 @@
* express or implied. See the License for the specific language governing * express or implied. See the License for the specific language governing
* permissions and limitations under the License. * permissions and limitations under the License.
*/ */
package software.amazon.kinesis.leases; package software.amazon.kinesis.leases.dynamodb;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
@ -40,33 +40,27 @@ import com.amazonaws.services.dynamodbv2.model.ScanRequest;
import com.amazonaws.services.dynamodbv2.model.ScanResult; import com.amazonaws.services.dynamodbv2.model.ScanResult;
import com.amazonaws.services.dynamodbv2.model.TableStatus; import com.amazonaws.services.dynamodbv2.model.TableStatus;
import com.amazonaws.services.dynamodbv2.model.UpdateItemRequest; import com.amazonaws.services.dynamodbv2.model.UpdateItemRequest;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.LeaseSerializer;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException; import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
import lombok.extern.slf4j.Slf4j;
/** /**
* An implementation of ILeaseManager that uses DynamoDB. * An implementation of {@link LeaseRefresher} that uses DynamoDB.
*/ */
@Slf4j @Slf4j
public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> { public class DynamoDBLeaseRefresher implements LeaseRefresher {
protected String table; protected String table;
protected AmazonDynamoDB dynamoDBClient; protected AmazonDynamoDB dynamoDBClient;
protected LeaseSerializer<T> serializer; protected LeaseSerializer serializer;
protected boolean consistentReads; protected boolean consistentReads;
/**
* Constructor.
*
* @param table leases table
* @param dynamoDBClient DynamoDB client to use
* @param serializer LeaseSerializer to use to convert to/from DynamoDB objects.
*/
public DynamoDBLeaseManager(String table, AmazonDynamoDB dynamoDBClient, LeaseSerializer<T> serializer) {
this(table, dynamoDBClient, serializer, false);
}
/** /**
* Constructor for test cases - allows control of consistent reads. Consistent reads should only be used for testing * Constructor for test cases - allows control of consistent reads. Consistent reads should only be used for testing
* - our code is meant to be resilient to inconsistent reads. Using consistent reads during testing speeds up * - our code is meant to be resilient to inconsistent reads. Using consistent reads during testing speeds up
@ -78,11 +72,8 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
* @param serializer lease serializer to use * @param serializer lease serializer to use
* @param consistentReads true if we want consistent reads for testing purposes. * @param consistentReads true if we want consistent reads for testing purposes.
*/ */
public DynamoDBLeaseManager(String table, AmazonDynamoDB dynamoDBClient, LeaseSerializer<T> serializer, boolean consistentReads) { public DynamoDBLeaseRefresher(@NonNull final String table, @NonNull final AmazonDynamoDB dynamoDBClient,
verifyNotNull(table, "Table name cannot be null"); @NonNull final LeaseSerializer serializer, boolean consistentReads) {
verifyNotNull(dynamoDBClient, "dynamoDBClient cannot be null");
verifyNotNull(serializer, "ILeaseSerializer cannot be null");
this.table = table; this.table = table;
this.dynamoDBClient = dynamoDBClient; this.dynamoDBClient = dynamoDBClient;
this.consistentReads = consistentReads; this.consistentReads = consistentReads;
@ -93,11 +84,8 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public boolean createLeaseTableIfNotExists(Long readCapacity, Long writeCapacity) public boolean createLeaseTableIfNotExists(@NonNull final Long readCapacity, @NonNull final Long writeCapacity)
throws ProvisionedThroughputException, DependencyException { throws ProvisionedThroughputException, DependencyException {
verifyNotNull(readCapacity, "readCapacity cannot be null");
verifyNotNull(writeCapacity, "writeCapacity cannot be null");
try { try {
if (tableStatus() != null) { if (tableStatus() != null) {
return false; return false;
@ -204,7 +192,7 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public List<T> listLeases() throws DependencyException, InvalidStateException, ProvisionedThroughputException { public List<Lease> listLeases() throws DependencyException, InvalidStateException, ProvisionedThroughputException {
return list(null); return list(null);
} }
@ -212,7 +200,8 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public boolean isLeaseTableEmpty() throws DependencyException, InvalidStateException, ProvisionedThroughputException { public boolean isLeaseTableEmpty()
throws DependencyException, InvalidStateException, ProvisionedThroughputException {
return list(1).isEmpty(); return list(1).isEmpty();
} }
@ -225,7 +214,7 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
* @throws DependencyException if DynamoDB scan fail in an unexpected way * @throws DependencyException if DynamoDB scan fail in an unexpected way
* @throws ProvisionedThroughputException if DynamoDB scan fail due to exceeded capacity * @throws ProvisionedThroughputException if DynamoDB scan fail due to exceeded capacity
*/ */
List<T> list(Integer limit) throws DependencyException, InvalidStateException, ProvisionedThroughputException { List<Lease> list(Integer limit) throws DependencyException, InvalidStateException, ProvisionedThroughputException {
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Listing leases from table {}", table); log.debug("Listing leases from table {}", table);
} }
@ -238,7 +227,7 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
try { try {
ScanResult scanResult = dynamoDBClient.scan(scanRequest); ScanResult scanResult = dynamoDBClient.scan(scanRequest);
List<T> result = new ArrayList<T>(); List<Lease> result = new ArrayList<>();
while (scanResult != null) { while (scanResult != null) {
for (Map<String, AttributeValue> item : scanResult.getItems()) { for (Map<String, AttributeValue> item : scanResult.getItems()) {
@ -286,10 +275,8 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public boolean createLeaseIfNotExists(T lease) public boolean createLeaseIfNotExists(@NonNull final Lease lease)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
verifyNotNull(lease, "lease cannot be null");
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Creating lease {}", lease); log.debug("Creating lease {}", lease);
} }
@ -308,7 +295,7 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
return false; return false;
} catch (AmazonClientException e) { } catch (AmazonClientException e) {
throw convertAndRethrowExceptions("create", lease.getLeaseKey(), e); throw convertAndRethrowExceptions("create", lease.leaseKey(), e);
} }
return true; return true;
@ -318,10 +305,8 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public T getLease(String leaseKey) public Lease getLease(@NonNull final String leaseKey)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
verifyNotNull(leaseKey, "leaseKey cannot be null");
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Getting lease with key {}", leaseKey); log.debug("Getting lease with key {}", leaseKey);
} }
@ -342,7 +327,7 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
return null; return null;
} else { } else {
T lease = serializer.fromDynamoRecord(dynamoRecord); final Lease lease = serializer.fromDynamoRecord(dynamoRecord);
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Got lease {}", lease); log.debug("Got lease {}", lease);
} }
@ -358,12 +343,10 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public boolean renewLease(T lease) public boolean renewLease(@NonNull final Lease lease)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
verifyNotNull(lease, "lease cannot be null");
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Renewing lease with key {}", lease.getLeaseKey()); log.debug("Renewing lease with key {}", lease.leaseKey());
} }
UpdateItemRequest request = new UpdateItemRequest(); UpdateItemRequest request = new UpdateItemRequest();
@ -377,26 +360,26 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
} catch (ConditionalCheckFailedException e) { } catch (ConditionalCheckFailedException e) {
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Lease renewal failed for lease with key {} because the lease counter was not {}", log.debug("Lease renewal failed for lease with key {} because the lease counter was not {}",
lease.getLeaseKey(), lease.getLeaseCounter()); lease.leaseKey(), lease.leaseCounter());
} }
// If we had a spurious retry during the Dynamo update, then this conditional PUT failure // If we had a spurious retry during the Dynamo update, then this conditional PUT failure
// might be incorrect. So, we get the item straight away and check if the lease owner + lease counter // might be incorrect. So, we get the item straight away and check if the lease owner + lease counter
// are what we expected. // are what we expected.
String expectedOwner = lease.getLeaseOwner(); String expectedOwner = lease.leaseOwner();
Long expectedCounter = lease.getLeaseCounter() + 1; Long expectedCounter = lease.leaseCounter() + 1;
T updatedLease = getLease(lease.getLeaseKey()); final Lease updatedLease = getLease(lease.leaseKey());
if (updatedLease == null || !expectedOwner.equals(updatedLease.getLeaseOwner()) || if (updatedLease == null || !expectedOwner.equals(updatedLease.leaseOwner()) ||
!expectedCounter.equals(updatedLease.getLeaseCounter())) { !expectedCounter.equals(updatedLease.leaseCounter())) {
return false; return false;
} }
log.info("Detected spurious renewal failure for lease with key {}, but recovered", lease.getLeaseKey()); log.info("Detected spurious renewal failure for lease with key {}, but recovered", lease.leaseKey());
} catch (AmazonClientException e) { } catch (AmazonClientException e) {
throw convertAndRethrowExceptions("renew", lease.getLeaseKey(), e); throw convertAndRethrowExceptions("renew", lease.leaseKey(), e);
} }
lease.setLeaseCounter(lease.getLeaseCounter() + 1); lease.leaseCounter(lease.leaseCounter() + 1);
return true; return true;
} }
@ -404,15 +387,14 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public boolean takeLease(T lease, String owner) public boolean takeLease(@NonNull final Lease lease, @NonNull final String owner)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
verifyNotNull(lease, "lease cannot be null"); final String oldOwner = lease.leaseOwner();
verifyNotNull(owner, "owner cannot be null");
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Taking lease with leaseKey {} from {} to {}", log.debug("Taking lease with leaseKey {} from {} to {}",
lease.getLeaseKey(), lease.leaseKey(),
lease.getLeaseOwner() == null ? "nobody" : lease.getLeaseOwner(), lease.leaseOwner() == null ? "nobody" : lease.leaseOwner(),
owner); owner);
} }
@ -430,16 +412,20 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
} catch (ConditionalCheckFailedException e) { } catch (ConditionalCheckFailedException e) {
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Lease renewal failed for lease with key {} because the lease counter was not {}", log.debug("Lease renewal failed for lease with key {} because the lease counter was not {}",
lease.getLeaseKey(), lease.getLeaseCounter()); lease.leaseKey(), lease.leaseCounter());
} }
return false; return false;
} catch (AmazonClientException e) { } catch (AmazonClientException e) {
throw convertAndRethrowExceptions("take", lease.getLeaseKey(), e); throw convertAndRethrowExceptions("take", lease.leaseKey(), e);
} }
lease.setLeaseCounter(lease.getLeaseCounter() + 1); lease.leaseCounter(lease.leaseCounter() + 1);
lease.setLeaseOwner(owner); lease.leaseOwner(owner);
if (oldOwner != null && !oldOwner.equals(owner)) {
lease.ownerSwitchesSinceCheckpoint(lease.ownerSwitchesSinceCheckpoint() + 1);
}
return true; return true;
} }
@ -448,14 +434,10 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public boolean evictLease(T lease) public boolean evictLease(@NonNull final Lease lease)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
verifyNotNull(lease, "lease cannot be null");
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Evicting lease with leaseKey {} owned by {}", log.debug("Evicting lease with leaseKey {} owned by {}", lease.leaseKey(), lease.leaseOwner());
lease.getLeaseKey(),
lease.getLeaseOwner());
} }
UpdateItemRequest request = new UpdateItemRequest(); UpdateItemRequest request = new UpdateItemRequest();
@ -472,16 +454,16 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
} catch (ConditionalCheckFailedException e) { } catch (ConditionalCheckFailedException e) {
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Lease eviction failed for lease with key {} because the lease owner was not {}", log.debug("Lease eviction failed for lease with key {} because the lease owner was not {}",
lease.getLeaseKey(), lease.getLeaseOwner()); lease.leaseKey(), lease.leaseOwner());
} }
return false; return false;
} catch (AmazonClientException e) { } catch (AmazonClientException e) {
throw convertAndRethrowExceptions("evict", lease.getLeaseKey(), e); throw convertAndRethrowExceptions("evict", lease.leaseKey(), e);
} }
lease.setLeaseOwner(null); lease.leaseOwner(null);
lease.setLeaseCounter(lease.getLeaseCounter() + 1); lease.leaseCounter(lease.leaseCounter() + 1);
return true; return true;
} }
@ -489,11 +471,11 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
public void deleteAll() throws DependencyException, InvalidStateException, ProvisionedThroughputException { public void deleteAll() throws DependencyException, InvalidStateException, ProvisionedThroughputException {
List<T> allLeases = listLeases(); List<Lease> allLeases = listLeases();
log.warn("Deleting {} items from table {}", allLeases.size(), table); log.warn("Deleting {} items from table {}", allLeases.size(), table);
for (T lease : allLeases) { for (final Lease lease : allLeases) {
DeleteItemRequest deleteRequest = new DeleteItemRequest(); DeleteItemRequest deleteRequest = new DeleteItemRequest();
deleteRequest.setTableName(table); deleteRequest.setTableName(table);
deleteRequest.setKey(serializer.getDynamoHashKey(lease)); deleteRequest.setKey(serializer.getDynamoHashKey(lease));
@ -506,11 +488,10 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public void deleteLease(T lease) throws DependencyException, InvalidStateException, ProvisionedThroughputException { public void deleteLease(@NonNull final Lease lease)
verifyNotNull(lease, "lease cannot be null"); throws DependencyException, InvalidStateException, ProvisionedThroughputException {
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Deleting lease with leaseKey {}", lease.getLeaseKey()); log.debug("Deleting lease with leaseKey {}", lease.leaseKey());
} }
DeleteItemRequest deleteRequest = new DeleteItemRequest(); DeleteItemRequest deleteRequest = new DeleteItemRequest();
@ -520,7 +501,7 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
try { try {
dynamoDBClient.deleteItem(deleteRequest); dynamoDBClient.deleteItem(deleteRequest);
} catch (AmazonClientException e) { } catch (AmazonClientException e) {
throw convertAndRethrowExceptions("delete", lease.getLeaseKey(), e); throw convertAndRethrowExceptions("delete", lease.leaseKey(), e);
} }
} }
@ -528,10 +509,8 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public boolean updateLease(T lease) public boolean updateLease(@NonNull final Lease lease)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
verifyNotNull(lease, "lease cannot be null");
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Updating lease {}", lease); log.debug("Updating lease {}", lease);
} }
@ -550,44 +529,53 @@ public class DynamoDBLeaseManager<T extends Lease> implements LeaseManager<T> {
} catch (ConditionalCheckFailedException e) { } catch (ConditionalCheckFailedException e) {
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Lease update failed for lease with key {} because the lease counter was not {}", log.debug("Lease update failed for lease with key {} because the lease counter was not {}",
lease.getLeaseKey(), lease.getLeaseCounter()); lease.leaseKey(), lease.leaseCounter());
} }
return false; return false;
} catch (AmazonClientException e) { } catch (AmazonClientException e) {
throw convertAndRethrowExceptions("update", lease.getLeaseKey(), e); throw convertAndRethrowExceptions("update", lease.leaseKey(), e);
} }
lease.setLeaseCounter(lease.getLeaseCounter() + 1); lease.leaseCounter(lease.leaseCounter() + 1);
return true; return true;
} }
/**
* {@inheritDoc}
*/
@Override
public ExtendedSequenceNumber getCheckpoint(String shardId)
throws ProvisionedThroughputException, InvalidStateException, DependencyException {
ExtendedSequenceNumber checkpoint = null;
Lease lease = getLease(shardId);
if (lease != null) {
checkpoint = lease.checkpoint();
}
return checkpoint;
}
/* /*
* This method contains boilerplate exception handling - it throws or returns something to be thrown. The * This method contains boilerplate exception handling - it throws or returns something to be thrown. The
* inconsistency there exists to satisfy the compiler when this method is used at the end of non-void methods. * inconsistency there exists to satisfy the compiler when this method is used at the end of non-void methods.
*/ */
protected DependencyException convertAndRethrowExceptions(String operation, String leaseKey, AmazonClientException e) protected DependencyException convertAndRethrowExceptions(String operation, String leaseKey,
AmazonClientException e)
throws ProvisionedThroughputException, InvalidStateException { throws ProvisionedThroughputException, InvalidStateException {
if (e instanceof ProvisionedThroughputExceededException) { if (e instanceof ProvisionedThroughputExceededException) {
log.warn("Provisioned Throughput on the lease table has been exceeded. It's recommended that you increase the IOPs on the table. Failure to increase the IOPs may cause the application to not make progress."); log.warn("Provisioned Throughput on the lease table has been exceeded. It's recommended that you increase" +
" the IOPs on the table. Failure to increase the IOPs may cause the application to not make" +
" progress.");
throw new ProvisionedThroughputException(e); throw new ProvisionedThroughputException(e);
} else if (e instanceof ResourceNotFoundException) { } else if (e instanceof ResourceNotFoundException) {
// @formatter:on // @formatter:on
throw new InvalidStateException(String.format("Cannot %s lease with key %s because table %s does not exist.", throw new InvalidStateException(
operation, String.format("Cannot %s lease with key %s because table %s does not exist.", operation, leaseKey,
leaseKey, table),
table),
e); e);
//@formatter:off //@formatter:off
} else { } else {
return new DependencyException(e); return new DependencyException(e);
} }
} }
private void verifyNotNull(Object object, String message) {
if (object == null) {
throw new IllegalArgumentException(message);
}
}
} }

View file

@ -12,7 +12,7 @@
* express or implied. See the License for the specific language governing * express or implied. See the License for the specific language governing
* permissions and limitations under the License. * permissions and limitations under the License.
*/ */
package software.amazon.kinesis.leases; package software.amazon.kinesis.leases.dynamodb;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.Collection; import java.util.Collection;
@ -30,25 +30,28 @@ import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeUnit;
import com.amazonaws.services.cloudwatch.model.StandardUnit; import com.amazonaws.services.cloudwatch.model.StandardUnit;
import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.LeaseRenewer;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException; import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
import software.amazon.kinesis.metrics.MetricsHelper;
import software.amazon.kinesis.metrics.ThreadSafeMetricsDelegatingScope;
import software.amazon.kinesis.metrics.IMetricsScope; import software.amazon.kinesis.metrics.IMetricsScope;
import software.amazon.kinesis.metrics.MetricsHelper;
import software.amazon.kinesis.metrics.MetricsLevel; import software.amazon.kinesis.metrics.MetricsLevel;
import software.amazon.kinesis.metrics.ThreadSafeMetricsDelegatingScope;
import lombok.extern.slf4j.Slf4j;
/** /**
* An implementation of ILeaseRenewer that uses DynamoDB via LeaseManager. * An implementation of {@link LeaseRenewer} that uses DynamoDB via {@link LeaseRefresher}.
*/ */
@Slf4j @Slf4j
public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> { public class DynamoDBLeaseRenewer implements LeaseRenewer {
private static final int RENEWAL_RETRIES = 2; private static final int RENEWAL_RETRIES = 2;
private final LeaseManager<T> leaseManager; private final LeaseRefresher leaseRefresher;
private final ConcurrentNavigableMap<String, T> ownedLeases = new ConcurrentSkipListMap<String, T>(); private final ConcurrentNavigableMap<String, Lease> ownedLeases = new ConcurrentSkipListMap<>();
private final String workerIdentifier; private final String workerIdentifier;
private final long leaseDurationNanos; private final long leaseDurationNanos;
private final ExecutorService executorService; private final ExecutorService executorService;
@ -56,14 +59,14 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
/** /**
* Constructor. * Constructor.
* *
* @param leaseManager LeaseManager to use * @param leaseRefresher LeaseRefresher to use
* @param workerIdentifier identifier of this worker * @param workerIdentifier identifier of this worker
* @param leaseDurationMillis duration of a lease in milliseconds * @param leaseDurationMillis duration of a lease in milliseconds
* @param executorService ExecutorService to use for renewing leases in parallel * @param executorService ExecutorService to use for renewing leases in parallel
*/ */
public DynamoDBLeaseRenewer(LeaseManager<T> leaseManager, String workerIdentifier, long leaseDurationMillis, public DynamoDBLeaseRenewer(final LeaseRefresher leaseRefresher, final String workerIdentifier,
ExecutorService executorService) { final long leaseDurationMillis, final ExecutorService executorService) {
this.leaseManager = leaseManager; this.leaseRefresher = leaseRefresher;
this.workerIdentifier = workerIdentifier; this.workerIdentifier = workerIdentifier;
this.leaseDurationNanos = TimeUnit.MILLISECONDS.toNanos(leaseDurationMillis); this.leaseDurationNanos = TimeUnit.MILLISECONDS.toNanos(leaseDurationMillis);
this.executorService = executorService; this.executorService = executorService;
@ -77,10 +80,7 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
// Due to the eventually consistent nature of ConcurrentNavigableMap iterators, this log entry may become // Due to the eventually consistent nature of ConcurrentNavigableMap iterators, this log entry may become
// inaccurate during iteration. // inaccurate during iteration.
log.debug("Worker {} holding %d leases: {}", log.debug("Worker {} holding {} leases: {}", workerIdentifier, ownedLeases.size(), ownedLeases);
workerIdentifier,
ownedLeases.size(),
ownedLeases);
} }
/* /*
@ -96,8 +96,8 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
* to getCurrentlyHeldLeases. They'll still cross paths, but they won't interleave their executions. * to getCurrentlyHeldLeases. They'll still cross paths, but they won't interleave their executions.
*/ */
int lostLeases = 0; int lostLeases = 0;
List<Future<Boolean>> renewLeaseTasks = new ArrayList<Future<Boolean>>(); List<Future<Boolean>> renewLeaseTasks = new ArrayList<>();
for (T lease : ownedLeases.descendingMap().values()) { for (Lease lease : ownedLeases.descendingMap().values()) {
renewLeaseTasks.add(executorService.submit(new RenewLeaseTask(lease, renewLeaseTaskMetricsScope))); renewLeaseTasks.add(executorService.submit(new RenewLeaseTask(lease, renewLeaseTaskMetricsScope)));
} }
int leasesInUnknownState = 0; int leasesInUnknownState = 0;
@ -132,10 +132,10 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
private class RenewLeaseTask implements Callable<Boolean> { private class RenewLeaseTask implements Callable<Boolean> {
private final T lease; private final Lease lease;
private final IMetricsScope metricsScope; private final IMetricsScope metricsScope;
public RenewLeaseTask(T lease, IMetricsScope metricsScope) { public RenewLeaseTask(Lease lease, IMetricsScope metricsScope) {
this.lease = lease; this.lease = lease;
this.metricsScope = metricsScope; this.metricsScope = metricsScope;
} }
@ -151,12 +151,12 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
} }
} }
private boolean renewLease(T lease) throws DependencyException, InvalidStateException { private boolean renewLease(Lease lease) throws DependencyException, InvalidStateException {
return renewLease(lease, false); return renewLease(lease, false);
} }
private boolean renewLease(T lease, boolean renewEvenIfExpired) throws DependencyException, InvalidStateException { private boolean renewLease(Lease lease, boolean renewEvenIfExpired) throws DependencyException, InvalidStateException {
String leaseKey = lease.getLeaseKey(); String leaseKey = lease.leaseKey();
boolean success = false; boolean success = false;
boolean renewedLease = false; boolean renewedLease = false;
@ -170,10 +170,10 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
// ShutdownException). // ShutdownException).
boolean isLeaseExpired = lease.isExpired(leaseDurationNanos, System.nanoTime()); boolean isLeaseExpired = lease.isExpired(leaseDurationNanos, System.nanoTime());
if (renewEvenIfExpired || !isLeaseExpired) { if (renewEvenIfExpired || !isLeaseExpired) {
renewedLease = leaseManager.renewLease(lease); renewedLease = leaseRefresher.renewLease(lease);
} }
if (renewedLease) { if (renewedLease) {
lease.setLastCounterIncrementNanos(System.nanoTime()); lease.lastCounterIncrementNanos(System.nanoTime());
} }
} }
@ -209,14 +209,14 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public Map<String, T> getCurrentlyHeldLeases() { public Map<String, Lease> getCurrentlyHeldLeases() {
Map<String, T> result = new HashMap<String, T>(); Map<String, Lease> result = new HashMap<>();
long now = System.nanoTime(); long now = System.nanoTime();
for (String leaseKey : ownedLeases.keySet()) { for (String leaseKey : ownedLeases.keySet()) {
T copy = getCopyOfHeldLease(leaseKey, now); Lease copy = getCopyOfHeldLease(leaseKey, now);
if (copy != null) { if (copy != null) {
result.put(copy.getLeaseKey(), copy); result.put(copy.leaseKey(), copy);
} }
} }
@ -227,7 +227,7 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public T getCurrentlyHeldLease(String leaseKey) { public Lease getCurrentlyHeldLease(String leaseKey) {
return getCopyOfHeldLease(leaseKey, System.nanoTime()); return getCopyOfHeldLease(leaseKey, System.nanoTime());
} }
@ -238,19 +238,19 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
* @param now current timestamp for old-ness checking * @param now current timestamp for old-ness checking
* @return non-authoritative copy of the held lease, or null if we don't currently hold it * @return non-authoritative copy of the held lease, or null if we don't currently hold it
*/ */
private T getCopyOfHeldLease(String leaseKey, long now) { private Lease getCopyOfHeldLease(String leaseKey, long now) {
T authoritativeLease = ownedLeases.get(leaseKey); Lease authoritativeLease = ownedLeases.get(leaseKey);
if (authoritativeLease == null) { if (authoritativeLease == null) {
return null; return null;
} else { } else {
T copy = null; Lease copy = null;
synchronized (authoritativeLease) { synchronized (authoritativeLease) {
copy = authoritativeLease.copy(); copy = authoritativeLease.copy();
} }
if (copy.isExpired(leaseDurationNanos, now)) { if (copy.isExpired(leaseDurationNanos, now)) {
log.info("getCurrentlyHeldLease not returning lease with key {} because it is expired", log.info("getCurrentlyHeldLease not returning lease with key {} because it is expired",
copy.getLeaseKey()); copy.leaseKey());
return null; return null;
} else { } else {
return copy; return copy;
@ -262,14 +262,14 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public boolean updateLease(T lease, UUID concurrencyToken) public boolean updateLease(Lease lease, UUID concurrencyToken)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
verifyNotNull(lease, "lease cannot be null"); verifyNotNull(lease, "lease cannot be null");
verifyNotNull(lease.getLeaseKey(), "leaseKey cannot be null"); verifyNotNull(lease.leaseKey(), "leaseKey cannot be null");
verifyNotNull(concurrencyToken, "concurrencyToken cannot be null"); verifyNotNull(concurrencyToken, "concurrencyToken cannot be null");
String leaseKey = lease.getLeaseKey(); String leaseKey = lease.leaseKey();
T authoritativeLease = ownedLeases.get(leaseKey); Lease authoritativeLease = ownedLeases.get(leaseKey);
if (authoritativeLease == null) { if (authoritativeLease == null) {
log.info("Worker {} could not update lease with key {} because it does not hold it", log.info("Worker {} could not update lease with key {} because it does not hold it",
@ -283,7 +283,7 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
* the lease was lost and regained between when the caller acquired his concurrency token and when the caller * the lease was lost and regained between when the caller acquired his concurrency token and when the caller
* called update. * called update.
*/ */
if (!authoritativeLease.getConcurrencyToken().equals(concurrencyToken)) { if (!authoritativeLease.concurrencyToken().equals(concurrencyToken)) {
log.info("Worker {} refusing to update lease with key {} because" log.info("Worker {} refusing to update lease with key {} because"
+ " concurrency tokens don't match", workerIdentifier, leaseKey); + " concurrency tokens don't match", workerIdentifier, leaseKey);
return false; return false;
@ -294,10 +294,10 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
try { try {
synchronized (authoritativeLease) { synchronized (authoritativeLease) {
authoritativeLease.update(lease); authoritativeLease.update(lease);
boolean updatedLease = leaseManager.updateLease(authoritativeLease); boolean updatedLease = leaseRefresher.updateLease(authoritativeLease);
if (updatedLease) { if (updatedLease) {
// Updates increment the counter // Updates increment the counter
authoritativeLease.setLastCounterIncrementNanos(System.nanoTime()); authoritativeLease.lastCounterIncrementNanos(System.nanoTime());
} else { } else {
/* /*
* If updateLease returns false, it means someone took the lease from us. Remove the lease * If updateLease returns false, it means someone took the lease from us. Remove the lease
@ -313,7 +313,7 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
* *
* 1) Concurrency token check passes * 1) Concurrency token check passes
* 2) Pause. Lose lease, re-acquire lease. This requires at least one lease counter update. * 2) Pause. Lose lease, re-acquire lease. This requires at least one lease counter update.
* 3) Unpause. leaseManager.updateLease fails conditional write due to counter updates, returns * 3) Unpause. leaseRefresher.updateLease fails conditional write due to counter updates, returns
* false. * false.
* 4) ownedLeases.remove(key, value) doesn't do anything because authoritativeLease does not * 4) ownedLeases.remove(key, value) doesn't do anything because authoritativeLease does not
* .equals() the re-acquired version in the map on the basis of lease counter. This is what we want. * .equals() the re-acquired version in the map on the basis of lease counter. This is what we want.
@ -337,24 +337,24 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public void addLeasesToRenew(Collection<T> newLeases) { public void addLeasesToRenew(Collection<Lease> newLeases) {
verifyNotNull(newLeases, "newLeases cannot be null"); verifyNotNull(newLeases, "newLeases cannot be null");
for (T lease : newLeases) { for (Lease lease : newLeases) {
if (lease.getLastCounterIncrementNanos() == null) { if (lease.lastCounterIncrementNanos() == null) {
log.info("addLeasesToRenew ignoring lease with key {} because it does not have lastRenewalNanos set", log.info("addLeasesToRenew ignoring lease with key {} because it does not have lastRenewalNanos set",
lease.getLeaseKey()); lease.leaseKey());
continue; continue;
} }
T authoritativeLease = lease.copy(); Lease authoritativeLease = lease.copy();
/* /*
* Assign a concurrency token when we add this to the set of currently owned leases. This ensures that * Assign a concurrency token when we add this to the set of currently owned leases. This ensures that
* every time we acquire a lease, it gets a new concurrency token. * every time we acquire a lease, it gets a new concurrency token.
*/ */
authoritativeLease.setConcurrencyToken(UUID.randomUUID()); authoritativeLease.concurrencyToken(UUID.randomUUID());
ownedLeases.put(authoritativeLease.getLeaseKey(), authoritativeLease); ownedLeases.put(authoritativeLease.leaseKey(), authoritativeLease);
} }
} }
@ -371,8 +371,8 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
* @param lease the lease to drop. * @param lease the lease to drop.
*/ */
@Override @Override
public void dropLease(T lease) { public void dropLease(Lease lease) {
ownedLeases.remove(lease.getLeaseKey()); ownedLeases.remove(lease.leaseKey());
} }
/** /**
@ -380,12 +380,12 @@ public class DynamoDBLeaseRenewer<T extends Lease> implements LeaseRenewer<T> {
*/ */
@Override @Override
public void initialize() throws DependencyException, InvalidStateException, ProvisionedThroughputException { public void initialize() throws DependencyException, InvalidStateException, ProvisionedThroughputException {
Collection<T> leases = leaseManager.listLeases(); Collection<Lease> leases = leaseRefresher.listLeases();
List<T> myLeases = new LinkedList<T>(); List<Lease> myLeases = new LinkedList<>();
boolean renewEvenIfExpired = true; boolean renewEvenIfExpired = true;
for (T lease : leases) { for (Lease lease : leases) {
if (workerIdentifier.equals(lease.getLeaseOwner())) { if (workerIdentifier.equals(lease.leaseOwner())) {
log.info(" Worker {} found lease {}", workerIdentifier, lease); log.info(" Worker {} found lease {}", workerIdentifier, lease);
// Okay to renew even if lease is expired, because we start with an empty list and we add the lease to // Okay to renew even if lease is expired, because we start with an empty list and we add the lease to
// our list only after a successful renew. So we don't need to worry about the edge case where we could // our list only after a successful renew. So we don't need to worry about the edge case where we could

View file

@ -0,0 +1,236 @@
/*
* Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Amazon Software License (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/asl/
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*/
package software.amazon.kinesis.leases.dynamodb;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import com.amazonaws.services.dynamodbv2.model.AttributeAction;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.AttributeValue;
import com.amazonaws.services.dynamodbv2.model.AttributeValueUpdate;
import com.amazonaws.services.dynamodbv2.model.ExpectedAttributeValue;
import com.amazonaws.services.dynamodbv2.model.KeySchemaElement;
import com.amazonaws.services.dynamodbv2.model.KeyType;
import com.amazonaws.services.dynamodbv2.model.ScalarAttributeType;
import com.google.common.base.Strings;
import software.amazon.kinesis.leases.DynamoUtils;
import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.LeaseSerializer;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
/**
* An implementation of ILeaseSerializer for basic Lease objects. Can also instantiate subclasses of Lease so that
* LeaseSerializer can be decorated by other classes if you need to add fields to leases.
*/
public class DynamoDBLeaseSerializer implements LeaseSerializer {
private static final String LEASE_KEY_KEY = "leaseKey";
private static final String LEASE_OWNER_KEY = "leaseOwner";
private static final String LEASE_COUNTER_KEY = "leaseCounter";
private static final String OWNER_SWITCHES_KEY = "ownerSwitchesSinceCheckpoint";
private static final String CHECKPOINT_SEQUENCE_NUMBER_KEY = "checkpoint";
private static final String CHECKPOINT_SUBSEQUENCE_NUMBER_KEY = "checkpointSubSequenceNumber";
private static final String PENDING_CHECKPOINT_SEQUENCE_KEY = "pendingCheckpoint";
private static final String PENDING_CHECKPOINT_SUBSEQUENCE_KEY = "pendingCheckpointSubSequenceNumber";
private static final String PARENT_SHARD_ID_KEY = "parentShardId";
@Override
public Map<String, AttributeValue> toDynamoRecord(final Lease lease) {
Map<String, AttributeValue> result = new HashMap<>();
result.put(LEASE_KEY_KEY, DynamoUtils.createAttributeValue(lease.leaseKey()));
result.put(LEASE_COUNTER_KEY, DynamoUtils.createAttributeValue(lease.leaseCounter()));
if (lease.leaseOwner() != null) {
result.put(LEASE_OWNER_KEY, DynamoUtils.createAttributeValue(lease.leaseOwner()));
}
result.put(OWNER_SWITCHES_KEY, DynamoUtils.createAttributeValue(lease.ownerSwitchesSinceCheckpoint()));
result.put(CHECKPOINT_SEQUENCE_NUMBER_KEY, DynamoUtils.createAttributeValue(lease.checkpoint().getSequenceNumber()));
result.put(CHECKPOINT_SUBSEQUENCE_NUMBER_KEY, DynamoUtils.createAttributeValue(lease.checkpoint().getSubSequenceNumber()));
if (lease.parentShardIds() != null && !lease.parentShardIds().isEmpty()) {
result.put(PARENT_SHARD_ID_KEY, DynamoUtils.createAttributeValue(lease.parentShardIds()));
}
if (lease.pendingCheckpoint() != null && !lease.pendingCheckpoint().getSequenceNumber().isEmpty()) {
result.put(PENDING_CHECKPOINT_SEQUENCE_KEY, DynamoUtils.createAttributeValue(lease.pendingCheckpoint().getSequenceNumber()));
result.put(PENDING_CHECKPOINT_SUBSEQUENCE_KEY, DynamoUtils.createAttributeValue(lease.pendingCheckpoint().getSubSequenceNumber()));
}
return result;
}
@Override
public Lease fromDynamoRecord(final Map<String, AttributeValue> dynamoRecord) {
Lease result = new Lease();
result.leaseKey(DynamoUtils.safeGetString(dynamoRecord, LEASE_KEY_KEY));
result.leaseOwner(DynamoUtils.safeGetString(dynamoRecord, LEASE_OWNER_KEY));
result.leaseCounter(DynamoUtils.safeGetLong(dynamoRecord, LEASE_COUNTER_KEY));
result.ownerSwitchesSinceCheckpoint(DynamoUtils.safeGetLong(dynamoRecord, OWNER_SWITCHES_KEY));
result.checkpoint(
new ExtendedSequenceNumber(
DynamoUtils.safeGetString(dynamoRecord, CHECKPOINT_SEQUENCE_NUMBER_KEY),
DynamoUtils.safeGetLong(dynamoRecord, CHECKPOINT_SUBSEQUENCE_NUMBER_KEY))
);
result.parentShardIds(DynamoUtils.safeGetSS(dynamoRecord, PARENT_SHARD_ID_KEY));
if (!Strings.isNullOrEmpty(DynamoUtils.safeGetString(dynamoRecord, PENDING_CHECKPOINT_SEQUENCE_KEY))) {
result.pendingCheckpoint(
new ExtendedSequenceNumber(
DynamoUtils.safeGetString(dynamoRecord, PENDING_CHECKPOINT_SEQUENCE_KEY),
DynamoUtils.safeGetLong(dynamoRecord, PENDING_CHECKPOINT_SUBSEQUENCE_KEY))
);
}
return result;
}
@Override
public Map<String, AttributeValue> getDynamoHashKey(final String leaseKey) {
Map<String, AttributeValue> result = new HashMap<>();
result.put(LEASE_KEY_KEY, DynamoUtils.createAttributeValue(leaseKey));
return result;
}
@Override
public Map<String, AttributeValue> getDynamoHashKey(final Lease lease) {
return getDynamoHashKey(lease.leaseKey());
}
@Override
public Map<String, ExpectedAttributeValue> getDynamoLeaseCounterExpectation(final Lease lease) {
return getDynamoLeaseCounterExpectation(lease.leaseCounter());
}
public Map<String, ExpectedAttributeValue> getDynamoLeaseCounterExpectation(final Long leaseCounter) {
Map<String, ExpectedAttributeValue> result = new HashMap<>();
ExpectedAttributeValue eav = new ExpectedAttributeValue(DynamoUtils.createAttributeValue(leaseCounter));
result.put(LEASE_COUNTER_KEY, eav);
return result;
}
@Override
public Map<String, ExpectedAttributeValue> getDynamoLeaseOwnerExpectation(final Lease lease) {
Map<String, ExpectedAttributeValue> result = new HashMap<>();
ExpectedAttributeValue eav;
if (lease.leaseOwner() == null) {
eav = new ExpectedAttributeValue(false);
} else {
eav = new ExpectedAttributeValue(DynamoUtils.createAttributeValue(lease.leaseOwner()));
}
result.put(LEASE_OWNER_KEY, eav);
return result;
}
@Override
public Map<String, ExpectedAttributeValue> getDynamoNonexistantExpectation() {
Map<String, ExpectedAttributeValue> result = new HashMap<>();
ExpectedAttributeValue expectedAV = new ExpectedAttributeValue(false);
result.put(LEASE_KEY_KEY, expectedAV);
return result;
}
@Override
public Map<String, AttributeValueUpdate> getDynamoLeaseCounterUpdate(final Lease lease) {
return getDynamoLeaseCounterUpdate(lease.leaseCounter());
}
public Map<String, AttributeValueUpdate> getDynamoLeaseCounterUpdate(Long leaseCounter) {
Map<String, AttributeValueUpdate> result = new HashMap<>();
AttributeValueUpdate avu =
new AttributeValueUpdate(DynamoUtils.createAttributeValue(leaseCounter + 1), AttributeAction.PUT);
result.put(LEASE_COUNTER_KEY, avu);
return result;
}
@Override
public Map<String, AttributeValueUpdate> getDynamoTakeLeaseUpdate(final Lease lease, String owner) {
Map<String, AttributeValueUpdate> result = new HashMap<>();
result.put(LEASE_OWNER_KEY, new AttributeValueUpdate(DynamoUtils.createAttributeValue(owner),
AttributeAction.PUT));
String oldOwner = lease.leaseOwner();
if (oldOwner != null && !oldOwner.equals(owner)) {
result.put(OWNER_SWITCHES_KEY, new AttributeValueUpdate(DynamoUtils.createAttributeValue(1L),
AttributeAction.ADD));
}
return result;
}
@Override
public Map<String, AttributeValueUpdate> getDynamoEvictLeaseUpdate(final Lease lease) {
Map<String, AttributeValueUpdate> result = new HashMap<>();
result.put(LEASE_OWNER_KEY, new AttributeValueUpdate(null, AttributeAction.DELETE));
return result;
}
@Override
public Map<String, AttributeValueUpdate> getDynamoUpdateLeaseUpdate(final Lease lease) {
Map<String, AttributeValueUpdate> result = new HashMap<>();
result.put(CHECKPOINT_SEQUENCE_NUMBER_KEY, new AttributeValueUpdate(DynamoUtils.createAttributeValue(lease.checkpoint().getSequenceNumber()),
AttributeAction.PUT));
result.put(CHECKPOINT_SUBSEQUENCE_NUMBER_KEY, new AttributeValueUpdate(DynamoUtils.createAttributeValue(lease.checkpoint().getSubSequenceNumber()),
AttributeAction.PUT));
result.put(OWNER_SWITCHES_KEY,
new AttributeValueUpdate(DynamoUtils.createAttributeValue(lease.ownerSwitchesSinceCheckpoint()),
AttributeAction.PUT));
if (lease.pendingCheckpoint() != null && !lease.pendingCheckpoint().getSequenceNumber().isEmpty()) {
result.put(PENDING_CHECKPOINT_SEQUENCE_KEY, new AttributeValueUpdate(DynamoUtils.createAttributeValue(lease.pendingCheckpoint().getSequenceNumber()), AttributeAction.PUT));
result.put(PENDING_CHECKPOINT_SUBSEQUENCE_KEY, new AttributeValueUpdate(DynamoUtils.createAttributeValue(lease.pendingCheckpoint().getSubSequenceNumber()), AttributeAction.PUT));
} else {
result.put(PENDING_CHECKPOINT_SEQUENCE_KEY, new AttributeValueUpdate().withAction(AttributeAction.DELETE));
result.put(PENDING_CHECKPOINT_SUBSEQUENCE_KEY, new AttributeValueUpdate().withAction(AttributeAction.DELETE));
}
return result;
}
@Override
public Collection<KeySchemaElement> getKeySchema() {
List<KeySchemaElement> keySchema = new ArrayList<>();
keySchema.add(new KeySchemaElement().withAttributeName(LEASE_KEY_KEY).withKeyType(KeyType.HASH));
return keySchema;
}
@Override
public Collection<AttributeDefinition> getAttributeDefinitions() {
List<AttributeDefinition> definitions = new ArrayList<>();
definitions.add(new AttributeDefinition().withAttributeName(LEASE_KEY_KEY)
.withAttributeType(ScalarAttributeType.S));
return definitions;
}
}

View file

@ -12,7 +12,7 @@
* express or implied. See the License for the specific language governing * express or implied. See the License for the specific language governing
* permissions and limitations under the License. * permissions and limitations under the License.
*/ */
package software.amazon.kinesis.leases; package software.amazon.kinesis.leases.dynamodb;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.Collection; import java.util.Collection;
@ -27,20 +27,23 @@ import java.util.concurrent.Callable;
import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeUnit;
import com.amazonaws.services.cloudwatch.model.StandardUnit; import com.amazonaws.services.cloudwatch.model.StandardUnit;
import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.LeaseTaker;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException; import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
import software.amazon.kinesis.metrics.MetricsHelper;
import software.amazon.kinesis.metrics.IMetricsScope; import software.amazon.kinesis.metrics.IMetricsScope;
import software.amazon.kinesis.metrics.MetricsHelper;
import software.amazon.kinesis.metrics.MetricsLevel; import software.amazon.kinesis.metrics.MetricsLevel;
import lombok.extern.slf4j.Slf4j;
/** /**
* An implementation of ILeaseTaker that uses DynamoDB via LeaseManager. * An implementation of {@link LeaseTaker} that uses DynamoDB via {@link LeaseRefresher}.
*/ */
@Slf4j @Slf4j
public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> { public class DynamoDBLeaseTaker implements LeaseTaker {
private static final int TAKE_RETRIES = 3; private static final int TAKE_RETRIES = 3;
private static final int SCAN_RETRIES = 1; private static final int SCAN_RETRIES = 1;
@ -53,17 +56,17 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
} }
}; };
private final LeaseManager<T> leaseManager; private final LeaseRefresher leaseRefresher;
private final String workerIdentifier; private final String workerIdentifier;
private final Map<String, T> allLeases = new HashMap<String, T>(); private final Map<String, Lease> allLeases = new HashMap<>();
private final long leaseDurationNanos; private final long leaseDurationNanos;
private int maxLeasesForWorker = Integer.MAX_VALUE; private int maxLeasesForWorker = Integer.MAX_VALUE;
private int maxLeasesToStealAtOneTime = 1; private int maxLeasesToStealAtOneTime = 1;
private long lastScanTimeNanos = 0L; private long lastScanTimeNanos = 0L;
public DynamoDBLeaseTaker(LeaseManager<T> leaseManager, String workerIdentifier, long leaseDurationMillis) { public DynamoDBLeaseTaker(LeaseRefresher leaseRefresher, String workerIdentifier, long leaseDurationMillis) {
this.leaseManager = leaseManager; this.leaseRefresher = leaseRefresher;
this.workerIdentifier = workerIdentifier; this.workerIdentifier = workerIdentifier;
this.leaseDurationNanos = TimeUnit.MILLISECONDS.toNanos(leaseDurationMillis); this.leaseDurationNanos = TimeUnit.MILLISECONDS.toNanos(leaseDurationMillis);
} }
@ -81,7 +84,7 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
* @param maxLeasesForWorker Max leases this Worker can handle at a time * @param maxLeasesForWorker Max leases this Worker can handle at a time
* @return LeaseTaker * @return LeaseTaker
*/ */
public DynamoDBLeaseTaker<T> withMaxLeasesForWorker(int maxLeasesForWorker) { public DynamoDBLeaseTaker withMaxLeasesForWorker(int maxLeasesForWorker) {
if (maxLeasesForWorker <= 0) { if (maxLeasesForWorker <= 0) {
throw new IllegalArgumentException("maxLeasesForWorker should be >= 1"); throw new IllegalArgumentException("maxLeasesForWorker should be >= 1");
} }
@ -97,7 +100,7 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
* @param maxLeasesToStealAtOneTime Steal up to this many leases at one time (for load balancing) * @param maxLeasesToStealAtOneTime Steal up to this many leases at one time (for load balancing)
* @return LeaseTaker * @return LeaseTaker
*/ */
public DynamoDBLeaseTaker<T> withMaxLeasesToStealAtOneTime(int maxLeasesToStealAtOneTime) { public DynamoDBLeaseTaker withMaxLeasesToStealAtOneTime(int maxLeasesToStealAtOneTime) {
if (maxLeasesToStealAtOneTime <= 0) { if (maxLeasesToStealAtOneTime <= 0) {
throw new IllegalArgumentException("maxLeasesToStealAtOneTime should be >= 1"); throw new IllegalArgumentException("maxLeasesToStealAtOneTime should be >= 1");
} }
@ -109,7 +112,7 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
* {@inheritDoc} * {@inheritDoc}
*/ */
@Override @Override
public Map<String, T> takeLeases() throws DependencyException, InvalidStateException { public Map<String, Lease> takeLeases() throws DependencyException, InvalidStateException {
return takeLeases(SYSTEM_CLOCK_CALLABLE); return takeLeases(SYSTEM_CLOCK_CALLABLE);
} }
@ -125,10 +128,10 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
* @throws DependencyException * @throws DependencyException
* @throws InvalidStateException * @throws InvalidStateException
*/ */
synchronized Map<String, T> takeLeases(Callable<Long> timeProvider) synchronized Map<String, Lease> takeLeases(Callable<Long> timeProvider)
throws DependencyException, InvalidStateException { throws DependencyException, InvalidStateException {
// Key is leaseKey // Key is leaseKey
Map<String, T> takenLeases = new HashMap<String, T>(); Map<String, Lease> takenLeases = new HashMap<>();
long startTime = System.currentTimeMillis(); long startTime = System.currentTimeMillis();
boolean success = false; boolean success = false;
@ -159,21 +162,21 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
return takenLeases; return takenLeases;
} }
List<T> expiredLeases = getExpiredLeases(); List<Lease> expiredLeases = getExpiredLeases();
Set<T> leasesToTake = computeLeasesToTake(expiredLeases); Set<Lease> leasesToTake = computeLeasesToTake(expiredLeases);
Set<String> untakenLeaseKeys = new HashSet<String>(); Set<String> untakenLeaseKeys = new HashSet<>();
for (T lease : leasesToTake) { for (Lease lease : leasesToTake) {
String leaseKey = lease.getLeaseKey(); String leaseKey = lease.leaseKey();
startTime = System.currentTimeMillis(); startTime = System.currentTimeMillis();
success = false; success = false;
try { try {
for (int i = 1; i <= TAKE_RETRIES; i++) { for (int i = 1; i <= TAKE_RETRIES; i++) {
try { try {
if (leaseManager.takeLease(lease, workerIdentifier)) { if (leaseRefresher.takeLease(lease, workerIdentifier)) {
lease.setLastCounterIncrementNanos(System.nanoTime()); lease.lastCounterIncrementNanos(System.nanoTime());
takenLeases.put(leaseKey, lease); takenLeases.put(leaseKey, lease);
} else { } else {
untakenLeaseKeys.add(leaseKey); untakenLeaseKeys.add(leaseKey);
@ -247,7 +250,7 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
*/ */
private void updateAllLeases(Callable<Long> timeProvider) private void updateAllLeases(Callable<Long> timeProvider)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
List<T> freshList = leaseManager.listLeases(); List<Lease> freshList = leaseRefresher.listLeases();
try { try {
lastScanTimeNanos = timeProvider.call(); lastScanTimeNanos = timeProvider.call();
} catch (Exception e) { } catch (Exception e) {
@ -255,29 +258,29 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
} }
// This set will hold the lease keys not updated by the previous listLeases call. // This set will hold the lease keys not updated by the previous listLeases call.
Set<String> notUpdated = new HashSet<String>(allLeases.keySet()); Set<String> notUpdated = new HashSet<>(allLeases.keySet());
// Iterate over all leases, finding ones to try to acquire that haven't changed since the last iteration // Iterate over all leases, finding ones to try to acquire that haven't changed since the last iteration
for (T lease : freshList) { for (Lease lease : freshList) {
String leaseKey = lease.getLeaseKey(); String leaseKey = lease.leaseKey();
T oldLease = allLeases.get(leaseKey); Lease oldLease = allLeases.get(leaseKey);
allLeases.put(leaseKey, lease); allLeases.put(leaseKey, lease);
notUpdated.remove(leaseKey); notUpdated.remove(leaseKey);
if (oldLease != null) { if (oldLease != null) {
// If we've seen this lease before... // If we've seen this lease before...
if (oldLease.getLeaseCounter().equals(lease.getLeaseCounter())) { if (oldLease.leaseCounter().equals(lease.leaseCounter())) {
// ...and the counter hasn't changed, propagate the lastRenewalNanos time from the old lease // ...and the counter hasn't changed, propagate the lastRenewalNanos time from the old lease
lease.setLastCounterIncrementNanos(oldLease.getLastCounterIncrementNanos()); lease.lastCounterIncrementNanos(oldLease.lastCounterIncrementNanos());
} else { } else {
// ...and the counter has changed, set lastRenewalNanos to the time of the scan. // ...and the counter has changed, set lastRenewalNanos to the time of the scan.
lease.setLastCounterIncrementNanos(lastScanTimeNanos); lease.lastCounterIncrementNanos(lastScanTimeNanos);
} }
} else { } else {
if (lease.getLeaseOwner() == null) { if (lease.leaseOwner() == null) {
// if this new lease is unowned, it's never been renewed. // if this new lease is unowned, it's never been renewed.
lease.setLastCounterIncrementNanos(0L); lease.lastCounterIncrementNanos(0L);
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Treating new lease with key {} as never renewed because it is new and unowned.", log.debug("Treating new lease with key {} as never renewed because it is new and unowned.",
@ -285,7 +288,7 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
} }
} else { } else {
// if this new lease is owned, treat it as renewed as of the scan // if this new lease is owned, treat it as renewed as of the scan
lease.setLastCounterIncrementNanos(lastScanTimeNanos); lease.lastCounterIncrementNanos(lastScanTimeNanos);
if (log.isDebugEnabled()) { if (log.isDebugEnabled()) {
log.debug("Treating new lease with key {} as recently renewed because it is new and owned.", log.debug("Treating new lease with key {} as recently renewed because it is new and owned.",
leaseKey); leaseKey);
@ -303,10 +306,10 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
/** /**
* @return list of leases that were expired as of our last scan. * @return list of leases that were expired as of our last scan.
*/ */
private List<T> getExpiredLeases() { private List<Lease> getExpiredLeases() {
List<T> expiredLeases = new ArrayList<T>(); List<Lease> expiredLeases = new ArrayList<>();
for (T lease : allLeases.values()) { for (Lease lease : allLeases.values()) {
if (lease.isExpired(leaseDurationNanos, lastScanTimeNanos)) { if (lease.isExpired(leaseDurationNanos, lastScanTimeNanos)) {
expiredLeases.add(lease); expiredLeases.add(lease);
} }
@ -318,13 +321,12 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
/** /**
* Compute the number of leases I should try to take based on the state of the system. * Compute the number of leases I should try to take based on the state of the system.
* *
* @param allLeases map of shardId to lease containing all leases
* @param expiredLeases list of leases we determined to be expired * @param expiredLeases list of leases we determined to be expired
* @return set of leases to take. * @return set of leases to take.
*/ */
private Set<T> computeLeasesToTake(List<T> expiredLeases) { private Set<Lease> computeLeasesToTake(List<Lease> expiredLeases) {
Map<String, Integer> leaseCounts = computeLeaseCounts(expiredLeases); Map<String, Integer> leaseCounts = computeLeaseCounts(expiredLeases);
Set<T> leasesToTake = new HashSet<T>(); Set<Lease> leasesToTake = new HashSet<>();
IMetricsScope metrics = MetricsHelper.getMetricsScope(); IMetricsScope metrics = MetricsHelper.getMetricsScope();
int numLeases = allLeases.size(); int numLeases = allLeases.size();
@ -383,13 +385,13 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
} }
} else { } else {
// If there are no expired leases and we need a lease, consider stealing. // If there are no expired leases and we need a lease, consider stealing.
List<T> leasesToSteal = chooseLeasesToSteal(leaseCounts, numLeasesToReachTarget, target); List<Lease> leasesToSteal = chooseLeasesToSteal(leaseCounts, numLeasesToReachTarget, target);
for (T leaseToSteal : leasesToSteal) { for (Lease leaseToSteal : leasesToSteal) {
log.info("Worker {} needed {} leases but none were expired, so it will steal lease {} from {}", log.info("Worker {} needed {} leases but none were expired, so it will steal lease {} from {}",
workerIdentifier, workerIdentifier,
numLeasesToReachTarget, numLeasesToReachTarget,
leaseToSteal.getLeaseKey(), leaseToSteal.leaseKey(),
leaseToSteal.getLeaseOwner()); leaseToSteal.leaseOwner());
leasesToTake.add(leaseToSteal); leasesToTake.add(leaseToSteal);
} }
} }
@ -428,8 +430,8 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
* @param target target # of leases per worker * @param target target # of leases per worker
* @return Leases to steal, or empty list if we should not steal * @return Leases to steal, or empty list if we should not steal
*/ */
private List<T> chooseLeasesToSteal(Map<String, Integer> leaseCounts, int needed, int target) { private List<Lease> chooseLeasesToSteal(Map<String, Integer> leaseCounts, int needed, int target) {
List<T> leasesToSteal = new ArrayList<>(); List<Lease> leasesToSteal = new ArrayList<>();
Entry<String, Integer> mostLoadedWorker = null; Entry<String, Integer> mostLoadedWorker = null;
// Find the most loaded worker // Find the most loaded worker
@ -476,10 +478,10 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
} }
String mostLoadedWorkerIdentifier = mostLoadedWorker.getKey(); String mostLoadedWorkerIdentifier = mostLoadedWorker.getKey();
List<T> candidates = new ArrayList<T>(); List<Lease> candidates = new ArrayList<>();
// Collect leases belonging to that worker // Collect leases belonging to that worker
for (T lease : allLeases.values()) { for (Lease lease : allLeases.values()) {
if (mostLoadedWorkerIdentifier.equals(lease.getLeaseOwner())) { if (mostLoadedWorkerIdentifier.equals(lease.leaseOwner())) {
candidates.add(lease); candidates.add(lease);
} }
} }
@ -499,13 +501,13 @@ public class DynamoDBLeaseTaker<T extends Lease> implements LeaseTaker<T> {
* @param expiredLeases list of leases that are currently expired * @param expiredLeases list of leases that are currently expired
* @return map of workerIdentifier to lease count * @return map of workerIdentifier to lease count
*/ */
private Map<String, Integer> computeLeaseCounts(List<T> expiredLeases) { private Map<String, Integer> computeLeaseCounts(List<Lease> expiredLeases) {
Map<String, Integer> leaseCounts = new HashMap<String, Integer>(); Map<String, Integer> leaseCounts = new HashMap<>();
// Compute the number of leases per worker by looking through allLeases and ignoring leases that have expired. // Compute the number of leases per worker by looking through allLeases and ignoring leases that have expired.
for (T lease : allLeases.values()) { for (Lease lease : allLeases.values()) {
if (!expiredLeases.contains(lease)) { if (!expiredLeases.contains(lease)) {
String leaseOwner = lease.getLeaseOwner(); String leaseOwner = lease.leaseOwner();
Integer oldCount = leaseCounts.get(leaseOwner); Integer oldCount = leaseCounts.get(leaseOwner);
if (oldCount == null) { if (oldCount == null) {
leaseCounts.put(leaseOwner, 1); leaseCounts.put(leaseOwner, 1);

View file

@ -20,8 +20,8 @@ import lombok.AccessLevel;
import lombok.NonNull; import lombok.NonNull;
import lombok.RequiredArgsConstructor; import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.leases.LeaseManager; import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.KinesisClientLease; import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.ShardInfo; import software.amazon.kinesis.leases.ShardInfo;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber; import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
@ -39,7 +39,7 @@ import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
public class BlockOnParentShardTask implements ITask { public class BlockOnParentShardTask implements ITask {
@NonNull @NonNull
private final ShardInfo shardInfo; private final ShardInfo shardInfo;
private final LeaseManager<KinesisClientLease> leaseManager; private final LeaseRefresher leaseRefresher;
// Sleep for this duration if the parent shards have not completed processing, or we encounter an exception. // Sleep for this duration if the parent shards have not completed processing, or we encounter an exception.
private final long parentShardPollIntervalMillis; private final long parentShardPollIntervalMillis;
@ -58,9 +58,9 @@ public class BlockOnParentShardTask implements ITask {
try { try {
boolean blockedOnParentShard = false; boolean blockedOnParentShard = false;
for (String shardId : shardInfo.parentShardIds()) { for (String shardId : shardInfo.parentShardIds()) {
KinesisClientLease lease = leaseManager.getLease(shardId); Lease lease = leaseRefresher.getLease(shardId);
if (lease != null) { if (lease != null) {
ExtendedSequenceNumber checkpoint = lease.getCheckpoint(); ExtendedSequenceNumber checkpoint = lease.checkpoint();
if ((checkpoint == null) || (!checkpoint.equals(ExtendedSequenceNumber.SHARD_END))) { if ((checkpoint == null) || (!checkpoint.equals(ExtendedSequenceNumber.SHARD_END))) {
log.debug("Shard {} is not yet done. Its current checkpoint is {}", shardId, checkpoint); log.debug("Shard {} is not yet done. Its current checkpoint is {}", shardId, checkpoint);
blockedOnParentShard = true; blockedOnParentShard = true;

View file

@ -191,7 +191,7 @@ class ConsumerStates {
@Override @Override
public ITask createTask(ShardConsumer consumer) { public ITask createTask(ShardConsumer consumer) {
return new BlockOnParentShardTask(consumer.shardInfo(), return new BlockOnParentShardTask(consumer.shardInfo(),
consumer.leaseManager(), consumer.leaseRefresher(),
consumer.parentShardPollIntervalMillis()); consumer.parentShardPollIntervalMillis());
} }
@ -320,7 +320,7 @@ class ConsumerStates {
consumer.recordProcessorCheckpointer(), consumer.recordProcessorCheckpointer(),
consumer.taskBackoffTimeMillis(), consumer.taskBackoffTimeMillis(),
consumer.skipShardSyncAtWorkerInitializationIfLeasesExist(), consumer.skipShardSyncAtWorkerInitializationIfLeasesExist(),
consumer.leaseManagerProxy(), consumer.shardDetector(),
throttlingReporter, throttlingReporter,
recordsFetcher.getRecords(), recordsFetcher.getRecords(),
consumer.shouldCallProcessRecordsEvenForEmptyRecordList(), consumer.shouldCallProcessRecordsEvenForEmptyRecordList(),
@ -528,14 +528,14 @@ class ConsumerStates {
public ITask createTask(ShardConsumer consumer) { public ITask createTask(ShardConsumer consumer) {
// TODO: set shutdown reason // TODO: set shutdown reason
return new ShutdownTask(consumer.shardInfo(), return new ShutdownTask(consumer.shardInfo(),
consumer.leaseManagerProxy(), consumer.shardDetector(),
consumer.recordProcessor(), consumer.recordProcessor(),
consumer.recordProcessorCheckpointer(), consumer.recordProcessorCheckpointer(),
consumer.shutdownReason(), consumer.shutdownReason(),
consumer.initialPositionInStream(), consumer.initialPositionInStream(),
consumer.cleanupLeasesOfCompletedShards(), consumer.cleanupLeasesOfCompletedShards(),
consumer.ignoreUnexpectedChildShards(), consumer.ignoreUnexpectedChildShards(),
consumer.leaseManager(), consumer.leaseRefresher(),
consumer.taskBackoffTimeMillis(), consumer.taskBackoffTimeMillis(),
consumer.getRecordsCache()); consumer.getRecordsCache());
} }

View file

@ -20,7 +20,7 @@ import lombok.NonNull;
import lombok.RequiredArgsConstructor; import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.checkpoint.RecordProcessorCheckpointer; import software.amazon.kinesis.checkpoint.RecordProcessorCheckpointer;
import software.amazon.kinesis.leases.LeaseManagerProxy; import software.amazon.kinesis.leases.ShardDetector;
import software.amazon.kinesis.leases.ShardInfo; import software.amazon.kinesis.leases.ShardInfo;
import software.amazon.kinesis.metrics.IMetricsScope; import software.amazon.kinesis.metrics.IMetricsScope;
import software.amazon.kinesis.metrics.MetricsHelper; import software.amazon.kinesis.metrics.MetricsHelper;
@ -77,7 +77,7 @@ public class ProcessTask implements ITask {
@NonNull final RecordProcessorCheckpointer recordProcessorCheckpointer, @NonNull final RecordProcessorCheckpointer recordProcessorCheckpointer,
final long backoffTimeMillis, final long backoffTimeMillis,
final boolean skipShardSyncAtWorkerInitializationIfLeasesExist, final boolean skipShardSyncAtWorkerInitializationIfLeasesExist,
final LeaseManagerProxy leaseManagerProxy, final ShardDetector shardDetector,
@NonNull final ThrottlingReporter throttlingReporter, @NonNull final ThrottlingReporter throttlingReporter,
final ProcessRecordsInput processRecordsInput, final ProcessRecordsInput processRecordsInput,
final boolean shouldCallProcessRecordsEvenForEmptyRecordList, final boolean shouldCallProcessRecordsEvenForEmptyRecordList,
@ -93,7 +93,7 @@ public class ProcessTask implements ITask {
Optional<Shard> currentShard = Optional.empty(); Optional<Shard> currentShard = Optional.empty();
if (!skipShardSyncAtWorkerInitializationIfLeasesExist) { if (!skipShardSyncAtWorkerInitializationIfLeasesExist) {
currentShard = leaseManagerProxy.listShards().stream() currentShard = shardDetector.listShards().stream()
.filter(shard -> shardInfo.shardId().equals(shard.getShardId())) .filter(shard -> shardInfo.shardId().equals(shard.getShardId()))
.findFirst(); .findFirst();
} }

View file

@ -35,9 +35,8 @@ import lombok.Synchronized;
import lombok.experimental.Accessors; import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.checkpoint.RecordProcessorCheckpointer; import software.amazon.kinesis.checkpoint.RecordProcessorCheckpointer;
import software.amazon.kinesis.leases.LeaseManager; import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.KinesisClientLease; import software.amazon.kinesis.leases.ShardDetector;
import software.amazon.kinesis.leases.LeaseManagerProxy;
import software.amazon.kinesis.leases.ShardInfo; import software.amazon.kinesis.leases.ShardInfo;
import software.amazon.kinesis.metrics.IMetricsFactory; import software.amazon.kinesis.metrics.IMetricsFactory;
import software.amazon.kinesis.metrics.MetricsCollectingTaskDecorator; import software.amazon.kinesis.metrics.MetricsCollectingTaskDecorator;
@ -60,7 +59,7 @@ public class ShardConsumer {
@NonNull @NonNull
private final String streamName; private final String streamName;
@NonNull @NonNull
private final LeaseManager<KinesisClientLease> leaseManager; private final LeaseRefresher leaseRefresher;
@NonNull @NonNull
private final ExecutorService executorService; private final ExecutorService executorService;
@NonNull @NonNull
@ -87,7 +86,7 @@ public class ShardConsumer {
private final boolean cleanupLeasesOfCompletedShards; private final boolean cleanupLeasesOfCompletedShards;
private final boolean ignoreUnexpectedChildShards; private final boolean ignoreUnexpectedChildShards;
@NonNull @NonNull
private final LeaseManagerProxy leaseManagerProxy; private final ShardDetector shardDetector;
@NonNull @NonNull
private final IMetricsFactory metricsFactory; private final IMetricsFactory metricsFactory;

View file

@ -16,9 +16,9 @@ package software.amazon.kinesis.lifecycle;
import java.util.concurrent.CountDownLatch; import java.util.concurrent.CountDownLatch;
import software.amazon.kinesis.processor.ShutdownNotificationAware; import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.KinesisClientLease;
import software.amazon.kinesis.leases.LeaseCoordinator; import software.amazon.kinesis.leases.LeaseCoordinator;
import software.amazon.kinesis.processor.ShutdownNotificationAware;
/** /**
* Contains callbacks for completion of stages in a requested record processor shutdown. * Contains callbacks for completion of stages in a requested record processor shutdown.
@ -26,8 +26,8 @@ import software.amazon.kinesis.leases.LeaseCoordinator;
*/ */
public class ShardConsumerShutdownNotification implements ShutdownNotification { public class ShardConsumerShutdownNotification implements ShutdownNotification {
private final LeaseCoordinator<KinesisClientLease> leaseCoordinator; private final LeaseCoordinator leaseCoordinator;
private final KinesisClientLease lease; private final Lease lease;
private final CountDownLatch shutdownCompleteLatch; private final CountDownLatch shutdownCompleteLatch;
private final CountDownLatch notificationCompleteLatch; private final CountDownLatch notificationCompleteLatch;
@ -48,8 +48,10 @@ public class ShardConsumerShutdownNotification implements ShutdownNotification {
* @param shutdownCompleteLatch * @param shutdownCompleteLatch
* used to inform the caller once the record processor is fully shutdown * used to inform the caller once the record processor is fully shutdown
*/ */
public ShardConsumerShutdownNotification(LeaseCoordinator<KinesisClientLease> leaseCoordinator, KinesisClientLease lease, public ShardConsumerShutdownNotification(final LeaseCoordinator leaseCoordinator,
CountDownLatch notificationCompleteLatch, CountDownLatch shutdownCompleteLatch) { final Lease lease,
final CountDownLatch notificationCompleteLatch,
final CountDownLatch shutdownCompleteLatch) {
this.leaseCoordinator = leaseCoordinator; this.leaseCoordinator = leaseCoordinator;
this.lease = lease; this.lease = lease;
this.notificationCompleteLatch = notificationCompleteLatch; this.notificationCompleteLatch = notificationCompleteLatch;

View file

@ -21,9 +21,8 @@ import lombok.NonNull;
import lombok.RequiredArgsConstructor; import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.checkpoint.RecordProcessorCheckpointer; import software.amazon.kinesis.checkpoint.RecordProcessorCheckpointer;
import software.amazon.kinesis.leases.LeaseManager; import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.KinesisClientLease; import software.amazon.kinesis.leases.ShardDetector;
import software.amazon.kinesis.leases.LeaseManagerProxy;
import software.amazon.kinesis.leases.ShardInfo; import software.amazon.kinesis.leases.ShardInfo;
import software.amazon.kinesis.leases.ShardSyncer; import software.amazon.kinesis.leases.ShardSyncer;
import software.amazon.kinesis.metrics.MetricsHelper; import software.amazon.kinesis.metrics.MetricsHelper;
@ -43,7 +42,7 @@ public class ShutdownTask implements ITask {
@NonNull @NonNull
private final ShardInfo shardInfo; private final ShardInfo shardInfo;
@NonNull @NonNull
private final LeaseManagerProxy leaseManagerProxy; private final ShardDetector shardDetector;
@NonNull @NonNull
private final RecordProcessor recordProcessor; private final RecordProcessor recordProcessor;
@NonNull @NonNull
@ -55,7 +54,7 @@ public class ShutdownTask implements ITask {
private final boolean cleanupLeasesOfCompletedShards; private final boolean cleanupLeasesOfCompletedShards;
private final boolean ignoreUnexpectedChildShards; private final boolean ignoreUnexpectedChildShards;
@NonNull @NonNull
private final LeaseManager<KinesisClientLease> leaseManager; private final LeaseRefresher leaseRefresher;
private final long backoffTimeMillis; private final long backoffTimeMillis;
@NonNull @NonNull
private final GetRecordsCache getRecordsCache; private final GetRecordsCache getRecordsCache;
@ -114,8 +113,8 @@ public class ShutdownTask implements ITask {
if (reason == ShutdownReason.TERMINATE) { if (reason == ShutdownReason.TERMINATE) {
log.debug("Looking for child shards of shard {}", shardInfo.shardId()); log.debug("Looking for child shards of shard {}", shardInfo.shardId());
// create leases for the child shards // create leases for the child shards
ShardSyncer.checkAndCreateLeasesForNewShards(leaseManagerProxy, ShardSyncer.checkAndCreateLeasesForNewShards(shardDetector,
leaseManager, leaseRefresher,
initialPositionInStream, initialPositionInStream,
cleanupLeasesOfCompletedShards, cleanupLeasesOfCompletedShards,
ignoreUnexpectedChildShards); ignoreUnexpectedChildShards);

View file

@ -95,7 +95,7 @@ public abstract class CheckpointImplTestBase {
} }
/** /**
* Test method to verify setCheckpoint and getCheckpoint methods. * Test method to verify checkpoint and checkpoint methods.
* *
* @throws Exception * @throws Exception
*/ */

View file

@ -108,7 +108,7 @@ public class InMemoryCheckpointer implements Checkpointer {
@Override @Override
public ExtendedSequenceNumber getCheckpoint(String shardId) throws KinesisClientLibException { public ExtendedSequenceNumber getCheckpoint(String shardId) throws KinesisClientLibException {
ExtendedSequenceNumber checkpoint = flushpoints.get(shardId); ExtendedSequenceNumber checkpoint = flushpoints.get(shardId);
log.debug("getCheckpoint shardId: {} checkpoint: {}", shardId, checkpoint); log.debug("checkpoint shardId: {} checkpoint: {}", shardId, checkpoint);
return checkpoint; return checkpoint;
} }

View file

@ -50,14 +50,12 @@ import com.amazonaws.services.kinesis.clientlibrary.exceptions.KinesisClientLibN
import software.amazon.kinesis.checkpoint.Checkpoint; import software.amazon.kinesis.checkpoint.Checkpoint;
import software.amazon.kinesis.checkpoint.CheckpointConfig; import software.amazon.kinesis.checkpoint.CheckpointConfig;
import software.amazon.kinesis.checkpoint.CheckpointFactory; import software.amazon.kinesis.checkpoint.CheckpointFactory;
import software.amazon.kinesis.leases.LeaseManager; import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseRefresher;
import software.amazon.kinesis.leases.KinesisClientLease;
import software.amazon.kinesis.leases.KinesisClientLibLeaseCoordinator;
import software.amazon.kinesis.leases.LeaseCoordinator; import software.amazon.kinesis.leases.LeaseCoordinator;
import software.amazon.kinesis.leases.LeaseManagementConfig; import software.amazon.kinesis.leases.LeaseManagementConfig;
import software.amazon.kinesis.leases.LeaseManagementFactory; import software.amazon.kinesis.leases.LeaseManagementFactory;
import software.amazon.kinesis.leases.DynamoDBLeaseManager; import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.LeaseManagerProxy; import software.amazon.kinesis.leases.ShardDetector;
import software.amazon.kinesis.leases.ShardInfo; import software.amazon.kinesis.leases.ShardInfo;
import software.amazon.kinesis.leases.ShardSyncTaskManager; import software.amazon.kinesis.leases.ShardSyncTaskManager;
import software.amazon.kinesis.lifecycle.InitializationInput; import software.amazon.kinesis.lifecycle.InitializationInput;
@ -69,9 +67,9 @@ import software.amazon.kinesis.lifecycle.ShutdownReason;
import software.amazon.kinesis.metrics.IMetricsFactory; import software.amazon.kinesis.metrics.IMetricsFactory;
import software.amazon.kinesis.metrics.MetricsConfig; import software.amazon.kinesis.metrics.MetricsConfig;
import software.amazon.kinesis.processor.Checkpointer; import software.amazon.kinesis.processor.Checkpointer;
import software.amazon.kinesis.processor.RecordProcessor;
import software.amazon.kinesis.processor.ProcessorConfig; import software.amazon.kinesis.processor.ProcessorConfig;
import software.amazon.kinesis.processor.ProcessorFactory; import software.amazon.kinesis.processor.ProcessorFactory;
import software.amazon.kinesis.processor.RecordProcessor;
import software.amazon.kinesis.retrieval.GetRecordsCache; import software.amazon.kinesis.retrieval.GetRecordsCache;
import software.amazon.kinesis.retrieval.RetrievalConfig; import software.amazon.kinesis.retrieval.RetrievalConfig;
import software.amazon.kinesis.retrieval.RetrievalFactory; import software.amazon.kinesis.retrieval.RetrievalFactory;
@ -107,13 +105,13 @@ public class SchedulerTest {
@Mock @Mock
private GetRecordsCache getRecordsCache; private GetRecordsCache getRecordsCache;
@Mock @Mock
private KinesisClientLibLeaseCoordinator leaseCoordinator; private LeaseCoordinator leaseCoordinator;
@Mock @Mock
private ShardSyncTaskManager shardSyncTaskManager; private ShardSyncTaskManager shardSyncTaskManager;
@Mock @Mock
private DynamoDBLeaseManager<KinesisClientLease> dynamoDBLeaseManager; private DynamoDBLeaseRefresher dynamoDBLeaseRefresher;
@Mock @Mock
private LeaseManagerProxy leaseManagerProxy; private ShardDetector shardDetector;
@Mock @Mock
private Checkpointer checkpoint; private Checkpointer checkpoint;
@ -131,8 +129,8 @@ public class SchedulerTest {
processorConfig = new ProcessorConfig(processorFactory); processorConfig = new ProcessorConfig(processorFactory);
retrievalConfig = new RetrievalConfig(streamName, amazonKinesis).retrievalFactory(retrievalFactory); retrievalConfig = new RetrievalConfig(streamName, amazonKinesis).retrievalFactory(retrievalFactory);
when(leaseCoordinator.leaseManager()).thenReturn(dynamoDBLeaseManager); when(leaseCoordinator.leaseRefresher()).thenReturn(dynamoDBLeaseRefresher);
when(shardSyncTaskManager.leaseManagerProxy()).thenReturn(leaseManagerProxy); when(shardSyncTaskManager.shardDetector()).thenReturn(shardDetector);
when(retrievalFactory.createGetRecordsCache(any(ShardInfo.class), any(IMetricsFactory.class))).thenReturn(getRecordsCache); when(retrievalFactory.createGetRecordsCache(any(ShardInfo.class), any(IMetricsFactory.class))).thenReturn(getRecordsCache);
scheduler = new Scheduler(checkpointConfig, coordinatorConfig, leaseManagementConfig, lifecycleConfig, scheduler = new Scheduler(checkpointConfig, coordinatorConfig, leaseManagementConfig, lifecycleConfig,
@ -236,11 +234,11 @@ public class SchedulerTest {
@Test @Test
public final void testInitializationFailureWithRetries() throws Exception { public final void testInitializationFailureWithRetries() throws Exception {
doNothing().when(leaseCoordinator).initialize(); doNothing().when(leaseCoordinator).initialize();
when(leaseManagerProxy.listShards()).thenThrow(new RuntimeException()); when(shardDetector.listShards()).thenThrow(new RuntimeException());
scheduler.run(); scheduler.run();
verify(leaseManagerProxy, times(Scheduler.MAX_INITIALIZATION_ATTEMPTS)).listShards(); verify(shardDetector, times(Scheduler.MAX_INITIALIZATION_ATTEMPTS)).listShards();
} }
@ -250,9 +248,9 @@ public class SchedulerTest {
final BigInteger startSeqNum = BigInteger.ONE; final BigInteger startSeqNum = BigInteger.ONE;
List<Shard> shardList = KinesisLocalFileDataCreator.createShardList(numShards, kinesisShardPrefix, startSeqNum); List<Shard> shardList = KinesisLocalFileDataCreator.createShardList(numShards, kinesisShardPrefix, startSeqNum);
Assert.assertEquals(numShards, shardList.size()); Assert.assertEquals(numShards, shardList.size());
List<KinesisClientLease> initialLeases = new ArrayList<KinesisClientLease>(); List<Lease> initialLeases = new ArrayList<Lease>();
for (Shard shard : shardList) { for (Shard shard : shardList) {
KinesisClientLease lease = ShardSyncer.newKCLLease(shard); Lease lease = ShardSyncer.newKCLLease(shard);
lease.setCheckpoint(ExtendedSequenceNumber.AT_TIMESTAMP); lease.setCheckpoint(ExtendedSequenceNumber.AT_TIMESTAMP);
initialLeases.add(lease); initialLeases.add(lease);
} }
@ -261,7 +259,7 @@ public class SchedulerTest {
private void runAndTestWorker(List<Shard> shardList, private void runAndTestWorker(List<Shard> shardList,
int threadPoolSize, int threadPoolSize,
List<KinesisClientLease> initialLeases, List<Lease> initialLeases,
int numberOfRecordsPerShard) throws Exception { int numberOfRecordsPerShard) throws Exception {
File file = KinesisLocalFileDataCreator.generateTempDataFile(shardList, numberOfRecordsPerShard, "unitTestWT001"); File file = KinesisLocalFileDataCreator.generateTempDataFile(shardList, numberOfRecordsPerShard, "unitTestWT001");
IKinesisProxy fileBasedProxy = new KinesisLocalFileProxy(file.getAbsolutePath()); IKinesisProxy fileBasedProxy = new KinesisLocalFileProxy(file.getAbsolutePath());
@ -288,7 +286,7 @@ public class SchedulerTest {
file.delete(); file.delete();
} }
private SchedulerThread runWorker(final List<KinesisClientLease> initialLeases) throws Exception { private SchedulerThread runWorker(final List<Lease> initialLeases) throws Exception {
final int maxRecords = 2; final int maxRecords = 2;
final long leaseDurationMillis = 10000L; final long leaseDurationMillis = 10000L;
@ -296,16 +294,16 @@ public class SchedulerTest {
final long idleTimeInMilliseconds = 2L; final long idleTimeInMilliseconds = 2L;
AmazonDynamoDB ddbClient = DynamoDBEmbedded.create().amazonDynamoDB(); AmazonDynamoDB ddbClient = DynamoDBEmbedded.create().amazonDynamoDB();
LeaseManager<KinesisClientLease> leaseManager = new KinesisClientLeaseManager("foo", ddbClient); LeaseManager<Lease> leaseRefresher = new LeaseManager("foo", ddbClient);
leaseManager.createLeaseTableIfNotExists(1L, 1L); leaseRefresher.createLeaseTableIfNotExists(1L, 1L);
for (KinesisClientLease initialLease : initialLeases) { for (Lease initialLease : initialLeases) {
leaseManager.createLeaseIfNotExists(initialLease); leaseRefresher.createLeaseIfNotExists(initialLease);
} }
checkpointConfig = new CheckpointConfig("foo", ddbClient, workerIdentifier) checkpointConfig = new CheckpointConfig("foo", ddbClient, workerIdentifier)
.failoverTimeMillis(leaseDurationMillis) .failoverTimeMillis(leaseDurationMillis)
.epsilonMillis(epsilonMillis) .epsilonMillis(epsilonMillis)
.leaseManager(leaseManager); .leaseRefresher(leaseRefresher);
leaseManagementConfig = new LeaseManagementConfig("foo", ddbClient, amazonKinesis, streamName, workerIdentifier) leaseManagementConfig = new LeaseManagementConfig("foo", ddbClient, amazonKinesis, streamName, workerIdentifier)
.failoverTimeMillis(leaseDurationMillis) .failoverTimeMillis(leaseDurationMillis)
.epsilonMillis(epsilonMillis); .epsilonMillis(epsilonMillis);
@ -323,7 +321,7 @@ public class SchedulerTest {
private void testWorker(List<Shard> shardList, private void testWorker(List<Shard> shardList,
int threadPoolSize, int threadPoolSize,
List<KinesisClientLease> initialLeases, List<Lease> initialLeases,
int numberOfRecordsPerShard, int numberOfRecordsPerShard,
IKinesisProxy kinesisProxy, IKinesisProxy kinesisProxy,
TestStreamletFactory recordProcessorFactory) throws Exception { TestStreamletFactory recordProcessorFactory) throws Exception {
@ -442,25 +440,20 @@ public class SchedulerTest {
} }
@Override @Override
public DynamoDBLeaseManager<KinesisClientLease> createLeaseManager() { public DynamoDBLeaseRefresher createLeaseRefresher() {
return dynamoDBLeaseManager; return dynamoDBLeaseRefresher;
} }
@Override @Override
public KinesisClientLibLeaseCoordinator createKinesisClientLibLeaseCoordinator() { public ShardDetector createShardDetector() {
return leaseCoordinator; return shardDetector;
}
@Override
public LeaseManagerProxy createLeaseManagerProxy() {
return leaseManagerProxy;
} }
} }
private class TestKinesisCheckpointFactory implements CheckpointFactory { private class TestKinesisCheckpointFactory implements CheckpointFactory {
@Override @Override
public Checkpointer createCheckpointer(final LeaseCoordinator<KinesisClientLease> leaseCoordinator, public Checkpointer createCheckpointer(final LeaseCoordinator leaseCoordinator,
final LeaseManager<KinesisClientLease> leaseManager) { final LeaseRefresher leaseRefresher) {
return checkpoint; return checkpoint;
} }
} }

View file

@ -50,7 +50,7 @@ public class WorkerTest {
@Mock @Mock
private KinesisClientLibLeaseCoordinator leaseCoordinator; private KinesisClientLibLeaseCoordinator leaseCoordinator;
@Mock @Mock
private ILeaseManager<KinesisClientLease> leaseManager; private ILeaseManager<KinesisClientLease> leaseRefresher;
@Mock @Mock
private software.amazon.kinesis.processor.IRecordProcessorFactory v1RecordProcessorFactory; private software.amazon.kinesis.processor.IRecordProcessorFactory v1RecordProcessorFactory;
@Mock @Mock
@ -151,7 +151,7 @@ public class WorkerTest {
final String dummyKinesisShardId = "kinesis-0-0"; final String dummyKinesisShardId = "kinesis-0-0";
ExecutorService execService = null; ExecutorService execService = null;
when(leaseCoordinator.leaseManager()).thenReturn(leaseManager); when(leaseCoordinator.leaseRefresher()).thenReturn(leaseRefresher);
Worker worker = Worker worker =
new Worker(stageName, new Worker(stageName,
@ -195,7 +195,7 @@ public class WorkerTest {
ExecutorService execService = null; ExecutorService execService = null;
when(leaseCoordinator.leaseManager()).thenReturn(leaseManager); when(leaseCoordinator.leaseRefresher()).thenReturn(leaseRefresher);
List<ShardInfo> initialState = createShardInfoList(ExtendedSequenceNumber.TRIM_HORIZON); List<ShardInfo> initialState = createShardInfoList(ExtendedSequenceNumber.TRIM_HORIZON);
List<ShardInfo> firstCheckpoint = createShardInfoList(new ExtendedSequenceNumber("1000")); List<ShardInfo> firstCheckpoint = createShardInfoList(new ExtendedSequenceNumber("1000"));
@ -270,7 +270,7 @@ public class WorkerTest {
final String dummyKinesisShardId = "kinesis-0-0"; final String dummyKinesisShardId = "kinesis-0-0";
final String anotherDummyKinesisShardId = "kinesis-0-1"; final String anotherDummyKinesisShardId = "kinesis-0-1";
ExecutorService execService = null; ExecutorService execService = null;
when(leaseCoordinator.leaseManager()).thenReturn(leaseManager); when(leaseCoordinator.leaseRefresher()).thenReturn(leaseRefresher);
Worker worker = Worker worker =
new Worker(stageName, new Worker(stageName,
@ -325,7 +325,7 @@ public class WorkerTest {
maxRecords, maxRecords,
idleTimeInMilliseconds, idleTimeInMilliseconds,
callProcessRecordsForEmptyRecordList, skipCheckpointValidationValue, INITIAL_POSITION_LATEST); callProcessRecordsForEmptyRecordList, skipCheckpointValidationValue, INITIAL_POSITION_LATEST);
when(leaseCoordinator.leaseManager()).thenReturn(leaseManager); when(leaseCoordinator.leaseRefresher()).thenReturn(leaseRefresher);
ExecutorService execService = Executors.newSingleThreadExecutor(); ExecutorService execService = Executors.newSingleThreadExecutor();
long shardPollInterval = 0L; long shardPollInterval = 0L;
Worker worker = Worker worker =
@ -392,7 +392,7 @@ public class WorkerTest {
List<Shard> shardList = createShardListWithOneSplit(); List<Shard> shardList = createShardListWithOneSplit();
List<KinesisClientLease> initialLeases = new ArrayList<KinesisClientLease>(); List<KinesisClientLease> initialLeases = new ArrayList<KinesisClientLease>();
KinesisClientLease lease = ShardSyncer.newKCLLease(shardList.get(0)); KinesisClientLease lease = ShardSyncer.newKCLLease(shardList.get(0));
lease.setCheckpoint(new ExtendedSequenceNumber("2")); lease.checkpoint(new ExtendedSequenceNumber("2"));
initialLeases.add(lease); initialLeases.add(lease);
runAndTestWorker(shardList, threadPoolSize, initialLeases, callProcessRecordsForEmptyRecordList, numberOfRecordsPerShard, config); runAndTestWorker(shardList, threadPoolSize, initialLeases, callProcessRecordsForEmptyRecordList, numberOfRecordsPerShard, config);
} }
@ -408,7 +408,7 @@ public class WorkerTest {
List<Shard> shardList = createShardListWithOneSplit(); List<Shard> shardList = createShardListWithOneSplit();
List<KinesisClientLease> initialLeases = new ArrayList<KinesisClientLease>(); List<KinesisClientLease> initialLeases = new ArrayList<KinesisClientLease>();
KinesisClientLease lease = ShardSyncer.newKCLLease(shardList.get(0)); KinesisClientLease lease = ShardSyncer.newKCLLease(shardList.get(0));
lease.setCheckpoint(new ExtendedSequenceNumber("2")); lease.checkpoint(new ExtendedSequenceNumber("2"));
initialLeases.add(lease); initialLeases.add(lease);
boolean callProcessRecordsForEmptyRecordList = true; boolean callProcessRecordsForEmptyRecordList = true;
RecordsFetcherFactory recordsFetcherFactory = new SimpleRecordsFetcherFactory(); RecordsFetcherFactory recordsFetcherFactory = new SimpleRecordsFetcherFactory();
@ -500,7 +500,7 @@ public class WorkerTest {
final List<KinesisClientLease> initialLeases = new ArrayList<KinesisClientLease>(); final List<KinesisClientLease> initialLeases = new ArrayList<KinesisClientLease>();
for (Shard shard : shardList) { for (Shard shard : shardList) {
KinesisClientLease lease = ShardSyncer.newKCLLease(shard); KinesisClientLease lease = ShardSyncer.newKCLLease(shard);
lease.setCheckpoint(ExtendedSequenceNumber.TRIM_HORIZON); lease.checkpoint(ExtendedSequenceNumber.TRIM_HORIZON);
initialLeases.add(lease); initialLeases.add(lease);
} }
@ -576,7 +576,7 @@ public class WorkerTest {
final List<KinesisClientLease> initialLeases = new ArrayList<KinesisClientLease>(); final List<KinesisClientLease> initialLeases = new ArrayList<KinesisClientLease>();
for (Shard shard : shardList) { for (Shard shard : shardList) {
KinesisClientLease lease = ShardSyncer.newKCLLease(shard); KinesisClientLease lease = ShardSyncer.newKCLLease(shard);
lease.setCheckpoint(ExtendedSequenceNumber.TRIM_HORIZON); lease.checkpoint(ExtendedSequenceNumber.TRIM_HORIZON);
initialLeases.add(lease); initialLeases.add(lease);
} }
@ -674,16 +674,16 @@ public class WorkerTest {
IMetricsFactory metricsFactory = mock(IMetricsFactory.class); IMetricsFactory metricsFactory = mock(IMetricsFactory.class);
ExtendedSequenceNumber checkpoint = new ExtendedSequenceNumber("123", 0L); ExtendedSequenceNumber checkpoint = new ExtendedSequenceNumber("123", 0L);
KinesisClientLeaseBuilder builder = new KinesisClientLeaseBuilder().withCheckpoint(checkpoint) KinesisClientLeaseBuilder builder = new KinesisClientLeaseBuilder().checkpoint(checkpoint)
.withConcurrencyToken(UUID.randomUUID()).withLastCounterIncrementNanos(0L).withLeaseCounter(0L) .concurrencyToken(UUID.randomUUID()).lastCounterIncrementNanos(0L).leaseCounter(0L)
.withOwnerSwitchesSinceCheckpoint(0L).withLeaseOwner("Self"); .ownerSwitchesSinceCheckpoint(0L).leaseOwner("Self");
final List<KinesisClientLease> leases = new ArrayList<>(); final List<KinesisClientLease> leases = new ArrayList<>();
final List<ShardInfo> currentAssignments = new ArrayList<>(); final List<ShardInfo> currentAssignments = new ArrayList<>();
KinesisClientLease lease = builder.withLeaseKey(String.format("shardId-%03d", 1)).build(); KinesisClientLease lease = builder.leaseKey(String.format("shardId-%03d", 1)).build();
leases.add(lease); leases.add(lease);
currentAssignments.add(new ShardInfo(lease.getLeaseKey(), lease.getConcurrencyToken().toString(), currentAssignments.add(new ShardInfo(lease.leaseKey(), lease.concurrencyToken().toString(),
lease.getParentShardIds(), lease.getCheckpoint())); lease.parentShardIds(), lease.checkpoint()));
when(leaseCoordinator.getAssignments()).thenAnswer(new Answer<List<KinesisClientLease>>() { when(leaseCoordinator.getAssignments()).thenAnswer(new Answer<List<KinesisClientLease>>() {
@ -762,16 +762,16 @@ public class WorkerTest {
IMetricsFactory metricsFactory = mock(IMetricsFactory.class); IMetricsFactory metricsFactory = mock(IMetricsFactory.class);
ExtendedSequenceNumber checkpoint = new ExtendedSequenceNumber("123", 0L); ExtendedSequenceNumber checkpoint = new ExtendedSequenceNumber("123", 0L);
KinesisClientLeaseBuilder builder = new KinesisClientLeaseBuilder().withCheckpoint(checkpoint) KinesisClientLeaseBuilder builder = new KinesisClientLeaseBuilder().checkpoint(checkpoint)
.withConcurrencyToken(UUID.randomUUID()).withLastCounterIncrementNanos(0L).withLeaseCounter(0L) .concurrencyToken(UUID.randomUUID()).lastCounterIncrementNanos(0L).leaseCounter(0L)
.withOwnerSwitchesSinceCheckpoint(0L).withLeaseOwner("Self"); .ownerSwitchesSinceCheckpoint(0L).leaseOwner("Self");
final List<KinesisClientLease> leases = new ArrayList<>(); final List<KinesisClientLease> leases = new ArrayList<>();
final List<ShardInfo> currentAssignments = new ArrayList<>(); final List<ShardInfo> currentAssignments = new ArrayList<>();
KinesisClientLease lease = builder.withLeaseKey(String.format("shardId-%03d", 1)).build(); KinesisClientLease lease = builder.leaseKey(String.format("shardId-%03d", 1)).build();
leases.add(lease); leases.add(lease);
currentAssignments.add(new ShardInfo(lease.getLeaseKey(), lease.getConcurrencyToken().toString(), currentAssignments.add(new ShardInfo(lease.leaseKey(), lease.concurrencyToken().toString(),
lease.getParentShardIds(), lease.getCheckpoint())); lease.parentShardIds(), lease.checkpoint()));
when(leaseCoordinator.getAssignments()).thenAnswer(new Answer<List<KinesisClientLease>>() { when(leaseCoordinator.getAssignments()).thenAnswer(new Answer<List<KinesisClientLease>>() {
@Override @Override
@ -827,16 +827,16 @@ public class WorkerTest {
IMetricsFactory metricsFactory = mock(IMetricsFactory.class); IMetricsFactory metricsFactory = mock(IMetricsFactory.class);
ExtendedSequenceNumber checkpoint = new ExtendedSequenceNumber("123", 0L); ExtendedSequenceNumber checkpoint = new ExtendedSequenceNumber("123", 0L);
KinesisClientLeaseBuilder builder = new KinesisClientLeaseBuilder().withCheckpoint(checkpoint) KinesisClientLeaseBuilder builder = new KinesisClientLeaseBuilder().checkpoint(checkpoint)
.withConcurrencyToken(UUID.randomUUID()).withLastCounterIncrementNanos(0L).withLeaseCounter(0L) .concurrencyToken(UUID.randomUUID()).lastCounterIncrementNanos(0L).leaseCounter(0L)
.withOwnerSwitchesSinceCheckpoint(0L).withLeaseOwner("Self"); .ownerSwitchesSinceCheckpoint(0L).leaseOwner("Self");
final List<KinesisClientLease> leases = new ArrayList<>(); final List<KinesisClientLease> leases = new ArrayList<>();
final List<ShardInfo> currentAssignments = new ArrayList<>(); final List<ShardInfo> currentAssignments = new ArrayList<>();
KinesisClientLease lease = builder.withLeaseKey(String.format("shardId-%03d", 1)).build(); KinesisClientLease lease = builder.leaseKey(String.format("shardId-%03d", 1)).build();
leases.add(lease); leases.add(lease);
currentAssignments.add(new ShardInfo(lease.getLeaseKey(), lease.getConcurrencyToken().toString(), currentAssignments.add(new ShardInfo(lease.leaseKey(), lease.concurrencyToken().toString(),
lease.getParentShardIds(), lease.getCheckpoint())); lease.parentShardIds(), lease.checkpoint()));
when(leaseCoordinator.getAssignments()).thenAnswer(new Answer<List<KinesisClientLease>>() { when(leaseCoordinator.getAssignments()).thenAnswer(new Answer<List<KinesisClientLease>>() {
@Override @Override
@ -1220,16 +1220,16 @@ public class WorkerTest {
IMetricsFactory metricsFactory = mock(IMetricsFactory.class); IMetricsFactory metricsFactory = mock(IMetricsFactory.class);
ExtendedSequenceNumber checkpoint = new ExtendedSequenceNumber("123", 0L); ExtendedSequenceNumber checkpoint = new ExtendedSequenceNumber("123", 0L);
KinesisClientLeaseBuilder builder = new KinesisClientLeaseBuilder().withCheckpoint(checkpoint) KinesisClientLeaseBuilder builder = new KinesisClientLeaseBuilder().checkpoint(checkpoint)
.withConcurrencyToken(UUID.randomUUID()).withLastCounterIncrementNanos(0L).withLeaseCounter(0L) .concurrencyToken(UUID.randomUUID()).lastCounterIncrementNanos(0L).leaseCounter(0L)
.withOwnerSwitchesSinceCheckpoint(0L).withLeaseOwner("Self"); .ownerSwitchesSinceCheckpoint(0L).leaseOwner("Self");
final List<KinesisClientLease> leases = new ArrayList<>(); final List<KinesisClientLease> leases = new ArrayList<>();
final List<ShardInfo> currentAssignments = new ArrayList<>(); final List<ShardInfo> currentAssignments = new ArrayList<>();
KinesisClientLease lease = builder.withLeaseKey(String.format("shardId-%03d", 1)).build(); KinesisClientLease lease = builder.leaseKey(String.format("shardId-%03d", 1)).build();
leases.add(lease); leases.add(lease);
currentAssignments.add(new ShardInfo(lease.getLeaseKey(), lease.getConcurrencyToken().toString(), currentAssignments.add(new ShardInfo(lease.leaseKey(), lease.concurrencyToken().toString(),
lease.getParentShardIds(), lease.getCheckpoint())); lease.parentShardIds(), lease.checkpoint()));
when(leaseCoordinator.getAssignments()).thenAnswer(new Answer<List<KinesisClientLease>>() { when(leaseCoordinator.getAssignments()).thenAnswer(new Answer<List<KinesisClientLease>>() {
@Override @Override
@ -1304,16 +1304,16 @@ public class WorkerTest {
IMetricsFactory metricsFactory = mock(IMetricsFactory.class); IMetricsFactory metricsFactory = mock(IMetricsFactory.class);
ExtendedSequenceNumber checkpoint = new ExtendedSequenceNumber("123", 0L); ExtendedSequenceNumber checkpoint = new ExtendedSequenceNumber("123", 0L);
KinesisClientLeaseBuilder builder = new KinesisClientLeaseBuilder().withCheckpoint(checkpoint) KinesisClientLeaseBuilder builder = new KinesisClientLeaseBuilder().checkpoint(checkpoint)
.withConcurrencyToken(UUID.randomUUID()).withLastCounterIncrementNanos(0L).withLeaseCounter(0L) .concurrencyToken(UUID.randomUUID()).lastCounterIncrementNanos(0L).leaseCounter(0L)
.withOwnerSwitchesSinceCheckpoint(0L).withLeaseOwner("Self"); .ownerSwitchesSinceCheckpoint(0L).leaseOwner("Self");
final List<KinesisClientLease> leases = new ArrayList<>(); final List<KinesisClientLease> leases = new ArrayList<>();
final List<ShardInfo> currentAssignments = new ArrayList<>(); final List<ShardInfo> currentAssignments = new ArrayList<>();
KinesisClientLease lease = builder.withLeaseKey(String.format("shardId-%03d", 1)).build(); KinesisClientLease lease = builder.leaseKey(String.format("shardId-%03d", 1)).build();
leases.add(lease); leases.add(lease);
currentAssignments.add(new ShardInfo(lease.getLeaseKey(), lease.getConcurrencyToken().toString(), currentAssignments.add(new ShardInfo(lease.leaseKey(), lease.concurrencyToken().toString(),
lease.getParentShardIds(), lease.getCheckpoint())); lease.parentShardIds(), lease.checkpoint()));
when(leaseCoordinator.getAssignments()).thenAnswer(new Answer<List<KinesisClientLease>>() { when(leaseCoordinator.getAssignments()).thenAnswer(new Answer<List<KinesisClientLease>>() {
@Override @Override
@ -1447,11 +1447,11 @@ public class WorkerTest {
final IRecordProcessor processor = mock(IRecordProcessor.class); final IRecordProcessor processor = mock(IRecordProcessor.class);
ExtendedSequenceNumber checkpoint = new ExtendedSequenceNumber("123", 0L); ExtendedSequenceNumber checkpoint = new ExtendedSequenceNumber("123", 0L);
KinesisClientLeaseBuilder builder = new KinesisClientLeaseBuilder().withCheckpoint(checkpoint) KinesisClientLeaseBuilder builder = new KinesisClientLeaseBuilder().checkpoint(checkpoint)
.withConcurrencyToken(UUID.randomUUID()).withLastCounterIncrementNanos(0L).withLeaseCounter(0L) .concurrencyToken(UUID.randomUUID()).lastCounterIncrementNanos(0L).leaseCounter(0L)
.withOwnerSwitchesSinceCheckpoint(0L).withLeaseOwner("Self"); .ownerSwitchesSinceCheckpoint(0L).leaseOwner("Self");
final List<KinesisClientLease> leases = new ArrayList<>(); final List<KinesisClientLease> leases = new ArrayList<>();
KinesisClientLease lease = builder.withLeaseKey(String.format("shardId-%03d", 1)).build(); KinesisClientLease lease = builder.leaseKey(String.format("shardId-%03d", 1)).build();
leases.add(lease); leases.add(lease);
doAnswer(new Answer<Boolean>() { doAnswer(new Answer<Boolean>() {
@ -1460,7 +1460,7 @@ public class WorkerTest {
workerInitialized.countDown(); workerInitialized.countDown();
return true; return true;
} }
}).when(leaseManager).waitUntilLeaseTableExists(anyLong(), anyLong()); }).when(leaseRefresher).waitUntilLeaseTableExists(anyLong(), anyLong());
doAnswer(new Answer<IRecordProcessor>() { doAnswer(new Answer<IRecordProcessor>() {
@Override @Override
public IRecordProcessor answer(InvocationOnMock invocation) throws Throwable { public IRecordProcessor answer(InvocationOnMock invocation) throws Throwable {
@ -1469,9 +1469,9 @@ public class WorkerTest {
} }
}).when(recordProcessorFactory).createProcessor(); }).when(recordProcessorFactory).createProcessor();
when(config.getWorkerIdentifier()).thenReturn("Self"); when(config.workerIdentifier()).thenReturn("Self");
when(leaseManager.listLeases()).thenReturn(leases); when(leaseRefresher.listLeases()).thenReturn(leases);
when(leaseManager.renewLease(leases.get(0))).thenReturn(true); when(leaseRefresher.renewLease(leases.get(0))).thenReturn(true);
when(executorService.submit(Matchers.<Callable<TaskResult>> any())) when(executorService.submit(Matchers.<Callable<TaskResult>> any()))
.thenAnswer(new ShutdownHandlingAnswer(taskFuture)); .thenAnswer(new ShutdownHandlingAnswer(taskFuture));
when(taskFuture.isDone()).thenReturn(true); when(taskFuture.isDone()).thenReturn(true);
@ -1481,7 +1481,7 @@ public class WorkerTest {
Worker worker = new Worker.Builder() Worker worker = new Worker.Builder()
.recordProcessorFactory(recordProcessorFactory) .recordProcessorFactory(recordProcessorFactory)
.config(config) .config(config)
.leaseManager(leaseManager) .leaseRefresher(leaseRefresher)
.kinesisProxy(kinesisProxy) .kinesisProxy(kinesisProxy)
.execService(executorService) .execService(executorService)
.workerStateChangeListener(workerStateChangeListener) .workerStateChangeListener(workerStateChangeListener)
@ -1513,7 +1513,7 @@ public class WorkerTest {
.config(config) .config(config)
.build(); .build();
Assert.assertNotNull(worker.getLeaseCoordinator().leaseManager()); Assert.assertNotNull(worker.getLeaseCoordinator().leaseRefresher());
} }
@SuppressWarnings("unchecked") @SuppressWarnings("unchecked")
@ -1521,14 +1521,14 @@ public class WorkerTest {
public void testBuilderWhenLeaseManagerIsSet() { public void testBuilderWhenLeaseManagerIsSet() {
IRecordProcessorFactory recordProcessorFactory = mock(IRecordProcessorFactory.class); IRecordProcessorFactory recordProcessorFactory = mock(IRecordProcessorFactory.class);
// Create an instance of ILeaseManager for injection and validation // Create an instance of ILeaseManager for injection and validation
ILeaseManager<KinesisClientLease> leaseManager = (ILeaseManager<KinesisClientLease>) mock(ILeaseManager.class); ILeaseManager<KinesisClientLease> leaseRefresher = (ILeaseManager<KinesisClientLease>) mock(ILeaseManager.class);
Worker worker = new Worker.Builder() Worker worker = new Worker.Builder()
.recordProcessorFactory(recordProcessorFactory) .recordProcessorFactory(recordProcessorFactory)
.config(config) .config(config)
.leaseManager(leaseManager) .leaseRefresher(leaseRefresher)
.build(); .build();
Assert.assertSame(leaseManager, worker.getLeaseCoordinator().leaseManager()); Assert.assertSame(leaseRefresher, worker.getLeaseCoordinator().leaseRefresher());
} }
private abstract class InjectableWorker extends Worker { private abstract class InjectableWorker extends Worker {
@ -1563,14 +1563,14 @@ public class WorkerTest {
} }
private KinesisClientLease makeLease(ExtendedSequenceNumber checkpoint, int shardId) { private KinesisClientLease makeLease(ExtendedSequenceNumber checkpoint, int shardId) {
return new KinesisClientLeaseBuilder().withCheckpoint(checkpoint).withConcurrencyToken(UUID.randomUUID()) return new KinesisClientLeaseBuilder().checkpoint(checkpoint).concurrencyToken(UUID.randomUUID())
.withLastCounterIncrementNanos(0L).withLeaseCounter(0L).withOwnerSwitchesSinceCheckpoint(0L) .lastCounterIncrementNanos(0L).leaseCounter(0L).ownerSwitchesSinceCheckpoint(0L)
.withLeaseOwner("Self").withLeaseKey(String.format("shardId-%03d", shardId)).build(); .leaseOwner("Self").leaseKey(String.format("shardId-%03d", shardId)).build();
} }
private ShardInfo makeShardInfo(KinesisClientLease lease) { private ShardInfo makeShardInfo(KinesisClientLease lease) {
return new ShardInfo(lease.getLeaseKey(), lease.getConcurrencyToken().toString(), lease.getParentShardIds(), return new ShardInfo(lease.leaseKey(), lease.concurrencyToken().toString(), lease.parentShardIds(),
lease.getCheckpoint()); lease.checkpoint());
} }
private static class ShutdownReasonMatcher extends TypeSafeDiagnosingMatcher<MetricsCollectingTaskDecorator> { private static class ShutdownReasonMatcher extends TypeSafeDiagnosingMatcher<MetricsCollectingTaskDecorator> {
@ -1785,7 +1785,7 @@ public class WorkerTest {
List<KinesisClientLease> initialLeases = new ArrayList<KinesisClientLease>(); List<KinesisClientLease> initialLeases = new ArrayList<KinesisClientLease>();
for (Shard shard : shardList) { for (Shard shard : shardList) {
KinesisClientLease lease = ShardSyncer.newKCLLease(shard); KinesisClientLease lease = ShardSyncer.newKCLLease(shard);
lease.setCheckpoint(ExtendedSequenceNumber.AT_TIMESTAMP); lease.checkpoint(ExtendedSequenceNumber.AT_TIMESTAMP);
initialLeases.add(lease); initialLeases.add(lease);
} }
runAndTestWorker(shardList, threadPoolSize, initialLeases, callProcessRecordsForEmptyRecordList, numberOfRecordsPerShard, config); runAndTestWorker(shardList, threadPoolSize, initialLeases, callProcessRecordsForEmptyRecordList, numberOfRecordsPerShard, config);
@ -1842,14 +1842,14 @@ public class WorkerTest {
final long idleTimeInMilliseconds = 2L; final long idleTimeInMilliseconds = 2L;
AmazonDynamoDB ddbClient = DynamoDBEmbedded.create().amazonDynamoDB(); AmazonDynamoDB ddbClient = DynamoDBEmbedded.create().amazonDynamoDB();
LeaseManager<KinesisClientLease> leaseManager = new KinesisClientLeaseManager("foo", ddbClient); LeaseManager<KinesisClientLease> leaseRefresher = new KinesisClientLeaseManager("foo", ddbClient);
leaseManager.createLeaseTableIfNotExists(1L, 1L); leaseRefresher.createLeaseTableIfNotExists(1L, 1L);
for (KinesisClientLease initialLease : initialLeases) { for (KinesisClientLease initialLease : initialLeases) {
leaseManager.createLeaseIfNotExists(initialLease); leaseRefresher.createLeaseIfNotExists(initialLease);
} }
KinesisClientLibLeaseCoordinator leaseCoordinator = KinesisClientLibLeaseCoordinator leaseCoordinator =
new KinesisClientLibLeaseCoordinator(leaseManager, new KinesisClientLibLeaseCoordinator(leaseRefresher,
stageName, stageName,
leaseDurationMillis, leaseDurationMillis,
epsilonMillis, epsilonMillis,

View file

@ -17,27 +17,27 @@ package software.amazon.kinesis.leases;
import java.util.Arrays; import java.util.Arrays;
import java.util.List; import java.util.List;
import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException; import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
import lombok.extern.slf4j.Slf4j;
/** /**
* Mock Lease Manager by randomly throwing Leasing Exceptions. * Mock LeaseRefresher by randomly throwing Leasing Exceptions.
* *
*/ */
@Slf4j @Slf4j
public class ExceptionThrowingLeaseManager implements LeaseManager<KinesisClientLease> { public class ExceptionThrowingLeaseRefresher implements LeaseRefresher {
private static final Throwable EXCEPTION_MSG = new Throwable("Test Exception"); private static final Throwable EXCEPTION_MSG = new Throwable("Test Exception");
// Use array below to control in what situations we want to throw exceptions. // Use array below to control in what situations we want to throw exceptions.
private int[] leaseManagerMethodCallingCount; private int[] leaseRefresherMethodCallingCount;
/** /**
* Methods which we support (simulate exceptions). * Methods which we support (simulate exceptions).
*/ */
public enum ExceptionThrowingLeaseManagerMethods { public enum ExceptionThrowingLeaseRefresherMethods {
CREATELEASETABLEIFNOTEXISTS(0), CREATELEASETABLEIFNOTEXISTS(0),
LEASETABLEEXISTS(1), LEASETABLEEXISTS(1),
WAITUNTILLEASETABLEEXISTS(2), WAITUNTILLEASETABLEEXISTS(2),
@ -54,7 +54,7 @@ public class ExceptionThrowingLeaseManager implements LeaseManager<KinesisClient
private Integer index; private Integer index;
ExceptionThrowingLeaseManagerMethods(Integer index) { ExceptionThrowingLeaseRefresherMethods(Integer index) {
this.index = index; this.index = index;
} }
@ -64,20 +64,20 @@ public class ExceptionThrowingLeaseManager implements LeaseManager<KinesisClient
} }
// Define which method should throw exception and when it should throw exception. // Define which method should throw exception and when it should throw exception.
private ExceptionThrowingLeaseManagerMethods methodThrowingException = ExceptionThrowingLeaseManagerMethods.NONE; private ExceptionThrowingLeaseRefresherMethods methodThrowingException = ExceptionThrowingLeaseRefresherMethods.NONE;
private int timeThrowingException = Integer.MAX_VALUE; private int timeThrowingException = Integer.MAX_VALUE;
// The real local lease manager which would do the real implementations. // The real local lease refresher which would do the real implementations.
private final LeaseManager<KinesisClientLease> leaseManager; private final LeaseRefresher leaseRefresher;
/** /**
* Constructor accepts lease manager as only argument. * Constructor accepts lease refresher as only argument.
* *
* @param leaseManager which will do the real implementations * @param leaseRefresher which will do the real implementations
*/ */
ExceptionThrowingLeaseManager(LeaseManager<KinesisClientLease> leaseManager) { ExceptionThrowingLeaseRefresher(LeaseRefresher leaseRefresher) {
this.leaseManager = leaseManager; this.leaseRefresher = leaseRefresher;
this.leaseManagerMethodCallingCount = new int[ExceptionThrowingLeaseManagerMethods.values().length]; this.leaseRefresherMethodCallingCount = new int[ExceptionThrowingLeaseRefresherMethods.values().length];
} }
/** /**
@ -86,7 +86,7 @@ public class ExceptionThrowingLeaseManager implements LeaseManager<KinesisClient
* @param method which would throw exception * @param method which would throw exception
* @param throwingTime defines what time to throw exception * @param throwingTime defines what time to throw exception
*/ */
void setLeaseLeaseManagerThrowingExceptionScenario(ExceptionThrowingLeaseManagerMethods method, int throwingTime) { void leaseRefresherThrowingExceptionScenario(ExceptionThrowingLeaseRefresherMethods method, int throwingTime) {
this.methodThrowingException = method; this.methodThrowingException = method;
this.timeThrowingException = throwingTime; this.timeThrowingException = throwingTime;
} }
@ -94,21 +94,21 @@ public class ExceptionThrowingLeaseManager implements LeaseManager<KinesisClient
/** /**
* Reset all parameters used for throwing exception. * Reset all parameters used for throwing exception.
*/ */
void clearLeaseManagerThrowingExceptionScenario() { void clearLeaseRefresherThrowingExceptionScenario() {
Arrays.fill(leaseManagerMethodCallingCount, 0); Arrays.fill(leaseRefresherMethodCallingCount, 0);
this.methodThrowingException = ExceptionThrowingLeaseManagerMethods.NONE; this.methodThrowingException = ExceptionThrowingLeaseRefresherMethods.NONE;
this.timeThrowingException = Integer.MAX_VALUE; this.timeThrowingException = Integer.MAX_VALUE;
} }
// Throw exception when the conditions are satisfied : // Throw exception when the conditions are satisfied :
// 1). method equals to methodThrowingException // 1). method equals to methodThrowingException
// 2). method calling count equals to what we want // 2). method calling count equals to what we want
private void throwExceptions(String methodName, ExceptionThrowingLeaseManagerMethods method) private void throwExceptions(String methodName, ExceptionThrowingLeaseRefresherMethods method)
throws DependencyException { throws DependencyException {
// Increase calling count for this method // Increase calling count for this method
leaseManagerMethodCallingCount[method.getIndex()]++; leaseRefresherMethodCallingCount[method.getIndex()]++;
if (method.equals(methodThrowingException) if (method.equals(methodThrowingException)
&& (leaseManagerMethodCallingCount[method.getIndex()] == timeThrowingException)) { && (leaseRefresherMethodCallingCount[method.getIndex()] == timeThrowingException)) {
// Throw Dependency Exception if all conditions are satisfied. // Throw Dependency Exception if all conditions are satisfied.
log.debug("Throwing DependencyException in {}", methodName); log.debug("Throwing DependencyException in {}", methodName);
throw new DependencyException(EXCEPTION_MSG); throw new DependencyException(EXCEPTION_MSG);
@ -119,94 +119,94 @@ public class ExceptionThrowingLeaseManager implements LeaseManager<KinesisClient
public boolean createLeaseTableIfNotExists(Long readCapacity, Long writeCapacity) public boolean createLeaseTableIfNotExists(Long readCapacity, Long writeCapacity)
throws ProvisionedThroughputException, DependencyException { throws ProvisionedThroughputException, DependencyException {
throwExceptions("createLeaseTableIfNotExists", throwExceptions("createLeaseTableIfNotExists",
ExceptionThrowingLeaseManagerMethods.CREATELEASETABLEIFNOTEXISTS); ExceptionThrowingLeaseRefresherMethods.CREATELEASETABLEIFNOTEXISTS);
return leaseManager.createLeaseTableIfNotExists(readCapacity, writeCapacity); return leaseRefresher.createLeaseTableIfNotExists(readCapacity, writeCapacity);
} }
@Override @Override
public boolean leaseTableExists() throws DependencyException { public boolean leaseTableExists() throws DependencyException {
throwExceptions("leaseTableExists", ExceptionThrowingLeaseManagerMethods.LEASETABLEEXISTS); throwExceptions("leaseTableExists", ExceptionThrowingLeaseRefresherMethods.LEASETABLEEXISTS);
return leaseManager.leaseTableExists(); return leaseRefresher.leaseTableExists();
} }
@Override @Override
public boolean waitUntilLeaseTableExists(long secondsBetweenPolls, long timeoutSeconds) throws DependencyException { public boolean waitUntilLeaseTableExists(long secondsBetweenPolls, long timeoutSeconds) throws DependencyException {
throwExceptions("waitUntilLeaseTableExists", ExceptionThrowingLeaseManagerMethods.WAITUNTILLEASETABLEEXISTS); throwExceptions("waitUntilLeaseTableExists", ExceptionThrowingLeaseRefresherMethods.WAITUNTILLEASETABLEEXISTS);
return leaseManager.waitUntilLeaseTableExists(secondsBetweenPolls, timeoutSeconds); return leaseRefresher.waitUntilLeaseTableExists(secondsBetweenPolls, timeoutSeconds);
} }
@Override @Override
public List<KinesisClientLease> listLeases() public List<Lease> listLeases()
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
throwExceptions("listLeases", ExceptionThrowingLeaseManagerMethods.LISTLEASES); throwExceptions("listLeases", ExceptionThrowingLeaseRefresherMethods.LISTLEASES);
return leaseManager.listLeases(); return leaseRefresher.listLeases();
} }
@Override @Override
public boolean createLeaseIfNotExists(KinesisClientLease lease) public boolean createLeaseIfNotExists(Lease lease)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
throwExceptions("createLeaseIfNotExists", ExceptionThrowingLeaseManagerMethods.CREATELEASEIFNOTEXISTS); throwExceptions("createLeaseIfNotExists", ExceptionThrowingLeaseRefresherMethods.CREATELEASEIFNOTEXISTS);
return leaseManager.createLeaseIfNotExists(lease); return leaseRefresher.createLeaseIfNotExists(lease);
} }
@Override @Override
public boolean renewLease(KinesisClientLease lease) public boolean renewLease(Lease lease)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
throwExceptions("renewLease", ExceptionThrowingLeaseManagerMethods.RENEWLEASE); throwExceptions("renewLease", ExceptionThrowingLeaseRefresherMethods.RENEWLEASE);
return leaseManager.renewLease(lease); return leaseRefresher.renewLease(lease);
} }
@Override @Override
public boolean takeLease(KinesisClientLease lease, String owner) public boolean takeLease(Lease lease, String owner)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
throwExceptions("takeLease", ExceptionThrowingLeaseManagerMethods.TAKELEASE); throwExceptions("takeLease", ExceptionThrowingLeaseRefresherMethods.TAKELEASE);
return leaseManager.takeLease(lease, owner); return leaseRefresher.takeLease(lease, owner);
} }
@Override @Override
public boolean evictLease(KinesisClientLease lease) public boolean evictLease(Lease lease)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
throwExceptions("evictLease", ExceptionThrowingLeaseManagerMethods.EVICTLEASE); throwExceptions("evictLease", ExceptionThrowingLeaseRefresherMethods.EVICTLEASE);
return leaseManager.evictLease(lease); return leaseRefresher.evictLease(lease);
} }
@Override @Override
public void deleteLease(KinesisClientLease lease) public void deleteLease(Lease lease)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
throwExceptions("deleteLease", ExceptionThrowingLeaseManagerMethods.DELETELEASE); throwExceptions("deleteLease", ExceptionThrowingLeaseRefresherMethods.DELETELEASE);
leaseManager.deleteLease(lease); leaseRefresher.deleteLease(lease);
} }
@Override @Override
public boolean updateLease(KinesisClientLease lease) public boolean updateLease(Lease lease)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
throwExceptions("updateLease", ExceptionThrowingLeaseManagerMethods.UPDATELEASE); throwExceptions("updateLease", ExceptionThrowingLeaseRefresherMethods.UPDATELEASE);
return leaseManager.updateLease(lease); return leaseRefresher.updateLease(lease);
} }
@Override @Override
public KinesisClientLease getLease(String shardId) public Lease getLease(String shardId)
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
throwExceptions("getLease", ExceptionThrowingLeaseManagerMethods.GETLEASE); throwExceptions("getLease", ExceptionThrowingLeaseRefresherMethods.GETLEASE);
return leaseManager.getLease(shardId); return leaseRefresher.getLease(shardId);
} }
@Override @Override
public void deleteAll() throws DependencyException, InvalidStateException, ProvisionedThroughputException { public void deleteAll() throws DependencyException, InvalidStateException, ProvisionedThroughputException {
throwExceptions("deleteAll", ExceptionThrowingLeaseManagerMethods.DELETEALL); throwExceptions("deleteAll", ExceptionThrowingLeaseRefresherMethods.DELETEALL);
leaseManager.deleteAll(); leaseRefresher.deleteAll();
} }
@Override @Override
@ -215,4 +215,9 @@ public class ExceptionThrowingLeaseManager implements LeaseManager<KinesisClient
return false; return false;
} }
@Override
public ExtendedSequenceNumber getCheckpoint(final String shardId)
throws ProvisionedThroughputException, InvalidStateException, DependencyException {
return null;
}
} }

View file

@ -1,83 +0,0 @@
/*
* Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Amazon Software License (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/asl/
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*/
package software.amazon.kinesis.leases;
import java.util.HashSet;
import java.util.Set;
import java.util.UUID;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
public class KinesisClientLeaseBuilder {
private String leaseKey;
private String leaseOwner;
private Long leaseCounter = 0L;
private UUID concurrencyToken;
private Long lastCounterIncrementNanos;
private ExtendedSequenceNumber checkpoint;
private ExtendedSequenceNumber pendingCheckpoint;
private Long ownerSwitchesSinceCheckpoint = 0L;
private Set<String> parentShardIds = new HashSet<>();
public KinesisClientLeaseBuilder withLeaseKey(String leaseKey) {
this.leaseKey = leaseKey;
return this;
}
public KinesisClientLeaseBuilder withLeaseOwner(String leaseOwner) {
this.leaseOwner = leaseOwner;
return this;
}
public KinesisClientLeaseBuilder withLeaseCounter(Long leaseCounter) {
this.leaseCounter = leaseCounter;
return this;
}
public KinesisClientLeaseBuilder withConcurrencyToken(UUID concurrencyToken) {
this.concurrencyToken = concurrencyToken;
return this;
}
public KinesisClientLeaseBuilder withLastCounterIncrementNanos(Long lastCounterIncrementNanos) {
this.lastCounterIncrementNanos = lastCounterIncrementNanos;
return this;
}
public KinesisClientLeaseBuilder withCheckpoint(ExtendedSequenceNumber checkpoint) {
this.checkpoint = checkpoint;
return this;
}
public KinesisClientLeaseBuilder withPendingCheckpoint(ExtendedSequenceNumber pendingCheckpoint) {
this.pendingCheckpoint = pendingCheckpoint;
return this;
}
public KinesisClientLeaseBuilder withOwnerSwitchesSinceCheckpoint(Long ownerSwitchesSinceCheckpoint) {
this.ownerSwitchesSinceCheckpoint = ownerSwitchesSinceCheckpoint;
return this;
}
public KinesisClientLeaseBuilder withParentShardIds(Set<String> parentShardIds) {
this.parentShardIds = parentShardIds;
return this;
}
public KinesisClientLease build() {
return new KinesisClientLease(leaseKey, leaseOwner, leaseCounter, concurrencyToken, lastCounterIncrementNanos,
checkpoint, pendingCheckpoint, ownerSwitchesSinceCheckpoint, parentShardIds);
}
}

View file

@ -0,0 +1,42 @@
/*
* Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Amazon Software License (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/asl/
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*/
package software.amazon.kinesis.leases;
import java.util.HashSet;
import java.util.Set;
import java.util.UUID;
import lombok.Setter;
import lombok.experimental.Accessors;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
@Setter
@Accessors(fluent = true)
public class LeaseBuilder {
private String leaseKey;
private String leaseOwner;
private Long leaseCounter = 0L;
private UUID concurrencyToken;
private Long lastCounterIncrementNanos;
private ExtendedSequenceNumber checkpoint;
private ExtendedSequenceNumber pendingCheckpoint;
private Long ownerSwitchesSinceCheckpoint = 0L;
private Set<String> parentShardIds = new HashSet<>();
public Lease build() {
return new Lease(leaseKey, leaseOwner, leaseCounter, concurrencyToken, lastCounterIncrementNanos,
checkpoint, pendingCheckpoint, ownerSwitchesSinceCheckpoint, parentShardIds);
}
}

View file

@ -30,17 +30,23 @@ import javax.swing.*;
import com.amazonaws.auth.AWSCredentialsProvider; import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient; import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseCoordinator;
import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseRefresher;
import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseSerializer;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.LeasingException; import software.amazon.kinesis.leases.exceptions.LeasingException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException; import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
import software.amazon.kinesis.metrics.CWMetricsFactory; import software.amazon.kinesis.metrics.CWMetricsFactory;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
import lombok.extern.slf4j.Slf4j;
@Slf4j @Slf4j
public class LeaseCoordinatorExerciser { public class LeaseCoordinatorExerciser {
private static final int MAX_LEASES_FOR_WORKER = Integer.MAX_VALUE;
private static final int MAX_LEASES_TO_STEAL_AT_ONE_TIME = 1;
private static final int MAX_LEASE_RENEWER_THREAD_COUNT = 20;
public static void main(String[] args) public static void main(String[] args)
throws InterruptedException, DependencyException, InvalidStateException, ProvisionedThroughputException, throws InterruptedException, DependencyException, InvalidStateException, ProvisionedThroughputException,
@ -55,38 +61,40 @@ public class LeaseCoordinatorExerciser {
new DefaultAWSCredentialsProviderChain(); new DefaultAWSCredentialsProviderChain();
AmazonDynamoDBClient ddb = new AmazonDynamoDBClient(creds); AmazonDynamoDBClient ddb = new AmazonDynamoDBClient(creds);
LeaseManager<KinesisClientLease> leaseManager = new KinesisClientDynamoDBLeaseManager("nagl_ShardProgress", ddb); LeaseRefresher leaseRefresher = new DynamoDBLeaseRefresher("nagl_ShardProgress", ddb, new DynamoDBLeaseSerializer(), true);
if (leaseManager.createLeaseTableIfNotExists(10L, 50L)) { if (leaseRefresher.createLeaseTableIfNotExists(10L, 50L)) {
log.info("Waiting for newly created lease table"); log.info("Waiting for newly created lease table");
if (!leaseManager.waitUntilLeaseTableExists(10, 300)) { if (!leaseRefresher.waitUntilLeaseTableExists(10, 300)) {
log.error("Table was not created in time"); log.error("Table was not created in time");
return; return;
} }
} }
CWMetricsFactory metricsFactory = new CWMetricsFactory(creds, "testNamespace", 30 * 1000, 1000); CWMetricsFactory metricsFactory = new CWMetricsFactory(creds, "testNamespace", 30 * 1000, 1000);
final List<LeaseCoordinator<KinesisClientLease>> coordinators = final List<LeaseCoordinator> coordinators = new ArrayList<>();
new ArrayList<LeaseCoordinator<KinesisClientLease>>();
for (int i = 0; i < numCoordinators; i++) { for (int i = 0; i < numCoordinators; i++) {
String workerIdentifier = "worker-" + Integer.toString(i); String workerIdentifier = "worker-" + Integer.toString(i);
LeaseCoordinator<KinesisClientLease> coord = new LeaseCoordinator<KinesisClientLease>(leaseManager, LeaseCoordinator coord = new DynamoDBLeaseCoordinator(leaseRefresher,
workerIdentifier, workerIdentifier,
leaseDurationMillis, leaseDurationMillis,
epsilonMillis, epsilonMillis,
MAX_LEASES_FOR_WORKER,
MAX_LEASES_TO_STEAL_AT_ONE_TIME,
MAX_LEASE_RENEWER_THREAD_COUNT,
metricsFactory); metricsFactory);
coordinators.add(coord); coordinators.add(coord);
} }
leaseManager.deleteAll(); leaseRefresher.deleteAll();
for (int i = 0; i < numLeases; i++) { for (int i = 0; i < numLeases; i++) {
KinesisClientLease lease = new KinesisClientLease(); Lease lease = new Lease();
lease.setLeaseKey(Integer.toString(i)); lease.leaseKey(Integer.toString(i));
lease.setCheckpoint(new ExtendedSequenceNumber("checkpoint")); lease.checkpoint(new ExtendedSequenceNumber("checkpoint"));
leaseManager.createLeaseIfNotExists(lease); leaseRefresher.createLeaseIfNotExists(lease);
} }
final JFrame frame = new JFrame("Test Visualizer"); final JFrame frame = new JFrame("Test Visualizer");
@ -97,10 +105,10 @@ public class LeaseCoordinatorExerciser {
frame.getContentPane().add(panel); frame.getContentPane().add(panel);
final Map<String, JLabel> labels = new HashMap<String, JLabel>(); final Map<String, JLabel> labels = new HashMap<String, JLabel>();
for (final LeaseCoordinator<KinesisClientLease> coord : coordinators) { for (final LeaseCoordinator coord : coordinators) {
JPanel coordPanel = new JPanel(); JPanel coordPanel = new JPanel();
coordPanel.setLayout(new BoxLayout(coordPanel, BoxLayout.X_AXIS)); coordPanel.setLayout(new BoxLayout(coordPanel, BoxLayout.X_AXIS));
final Button button = new Button("Stop " + coord.getWorkerIdentifier()); final Button button = new Button("Stop " + coord.workerIdentifier());
button.setMaximumSize(new Dimension(200, 50)); button.setMaximumSize(new Dimension(200, 50));
button.addActionListener(new ActionListener() { button.addActionListener(new ActionListener() {
@ -108,14 +116,14 @@ public class LeaseCoordinatorExerciser {
public void actionPerformed(ActionEvent arg0) { public void actionPerformed(ActionEvent arg0) {
if (coord.isRunning()) { if (coord.isRunning()) {
coord.stop(); coord.stop();
button.setLabel("Start " + coord.getWorkerIdentifier()); button.setLabel("Start " + coord.workerIdentifier());
} else { } else {
try { try {
coord.start(); coord.start();
} catch (LeasingException e) { } catch (LeasingException e) {
log.error("{}", e); log.error("{}", e);
} }
button.setLabel("Stop " + coord.getWorkerIdentifier()); button.setLabel("Stop " + coord.workerIdentifier());
} }
} }
@ -124,7 +132,7 @@ public class LeaseCoordinatorExerciser {
JLabel label = new JLabel(); JLabel label = new JLabel();
coordPanel.add(label); coordPanel.add(label);
labels.put(coord.getWorkerIdentifier(), label); labels.put(coord.workerIdentifier(), label);
panel.add(coordPanel); panel.add(coordPanel);
} }
@ -141,17 +149,17 @@ public class LeaseCoordinatorExerciser {
@Override @Override
public void run() { public void run() {
while (true) { while (true) {
for (LeaseCoordinator<KinesisClientLease> coord : coordinators) { for (LeaseCoordinator coord : coordinators) {
String workerIdentifier = coord.getWorkerIdentifier(); String workerIdentifier = coord.workerIdentifier();
JLabel label = labels.get(workerIdentifier); JLabel label = labels.get(workerIdentifier);
List<KinesisClientLease> asgn = new ArrayList<KinesisClientLease>(coord.getAssignments()); List<Lease> asgn = new ArrayList<>(coord.getAssignments());
Collections.sort(asgn, new Comparator<KinesisClientLease>() { Collections.sort(asgn, new Comparator<Lease>() {
@Override @Override
public int compare(KinesisClientLease arg0, KinesisClientLease arg1) { public int compare(final Lease arg0, final Lease arg1) {
return arg0.getLeaseKey().compareTo(arg1.getLeaseKey()); return arg0.leaseKey().compareTo(arg1.leaseKey());
} }
}); });
@ -160,19 +168,19 @@ public class LeaseCoordinatorExerciser {
builder.append("<html>"); builder.append("<html>");
builder.append(workerIdentifier).append(":").append(asgn.size()).append(" "); builder.append(workerIdentifier).append(":").append(asgn.size()).append(" ");
for (KinesisClientLease lease : asgn) { for (Lease lease : asgn) {
String leaseKey = lease.getLeaseKey(); String leaseKey = lease.leaseKey();
String lastOwner = lastOwners.get(leaseKey); String lastOwner = lastOwners.get(leaseKey);
// Color things green when they switch owners, decay the green-ness over time. // Color things green when they switch owners, decay the green-ness over time.
Integer greenNess = greenNesses.get(leaseKey); Integer greenNess = greenNesses.get(leaseKey);
if (greenNess == null || lastOwner == null || !lastOwner.equals(lease.getLeaseOwner())) { if (greenNess == null || lastOwner == null || !lastOwner.equals(lease.leaseOwner())) {
greenNess = 200; greenNess = 200;
} else { } else {
greenNess = Math.max(0, greenNess - 20); greenNess = Math.max(0, greenNess - 20);
} }
greenNesses.put(leaseKey, greenNess); greenNesses.put(leaseKey, greenNess);
lastOwners.put(leaseKey, lease.getLeaseOwner()); lastOwners.put(leaseKey, lease.leaseOwner());
builder.append(String.format("<font color=\"%s\">%03d</font>", builder.append(String.format("<font color=\"%s\">%03d</font>",
String.format("#00%02x00", greenNess), String.format("#00%02x00", greenNess),
@ -203,7 +211,7 @@ public class LeaseCoordinatorExerciser {
frame.pack(); frame.pack();
frame.setVisible(true); frame.setVisible(true);
for (LeaseCoordinator<KinesisClientLease> coord : coordinators) { for (LeaseCoordinator coord : coordinators) {
coord.start(); coord.start();
} }
} }

View file

@ -21,16 +21,19 @@ import org.junit.runner.Description;
import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient; import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient;
import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseRefresher;
import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseSerializer;
import software.amazon.kinesis.metrics.MetricsHelper; import software.amazon.kinesis.metrics.MetricsHelper;
import software.amazon.kinesis.metrics.NullMetricsFactory; import software.amazon.kinesis.metrics.NullMetricsFactory;
import lombok.extern.slf4j.Slf4j;
@Ignore
@Slf4j @Slf4j
@Ignore
public class LeaseIntegrationTest { public class LeaseIntegrationTest {
private LeaseSerializer leaseSerializer = new DynamoDBLeaseSerializer();
protected static KinesisClientDynamoDBLeaseManager leaseManager; protected static DynamoDBLeaseRefresher leaseRefresher;
protected static AmazonDynamoDBClient ddbClient = protected static AmazonDynamoDBClient ddbClient =
new AmazonDynamoDBClient(new DefaultAWSCredentialsProviderChain()); new AmazonDynamoDBClient(new DefaultAWSCredentialsProviderChain());
@ -39,25 +42,25 @@ public class LeaseIntegrationTest {
@Override @Override
protected void starting(Description description) { protected void starting(Description description) {
if (leaseManager == null) { if (leaseRefresher == null) {
// Do some static setup once per class. // Do some static setup once per class.
leaseManager = new KinesisClientDynamoDBLeaseManager("nagl_ShardProgress", ddbClient, true); leaseRefresher = new DynamoDBLeaseRefresher("nagl_ShardProgress", ddbClient, leaseSerializer, true);
MetricsHelper.startScope(new NullMetricsFactory()); MetricsHelper.startScope(new NullMetricsFactory());
} }
try { try {
if (!leaseManager.leaseTableExists()) { if (!leaseRefresher.leaseTableExists()) {
log.info("Creating lease table"); log.info("Creating lease table");
leaseManager.createLeaseTableIfNotExists(10L, 10L); leaseRefresher.createLeaseTableIfNotExists(10L, 10L);
leaseManager.waitUntilLeaseTableExists(10, 500); leaseRefresher.waitUntilLeaseTableExists(10, 500);
} }
log.info("Beginning test case {}", description.getMethodName()); log.info("Beginning test case {}", description.getMethodName());
for (KinesisClientLease lease : leaseManager.listLeases()) { for (Lease lease : leaseRefresher.listLeases()) {
leaseManager.deleteLease(lease); leaseRefresher.deleteLease(lease);
} }
} catch (Exception e) { } catch (Exception e) {
String message = String message =
@ -69,3 +72,4 @@ public class LeaseIntegrationTest {
}; };
} }

View file

@ -37,6 +37,8 @@ import com.amazonaws.services.kinesis.model.DescribeStreamSummaryRequest;
import com.amazonaws.services.kinesis.model.Shard; import com.amazonaws.services.kinesis.model.Shard;
import com.amazonaws.services.kinesis.model.StreamStatus; import com.amazonaws.services.kinesis.model.StreamStatus;
import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseRefresher;
import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseSerializer;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException; import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
@ -48,8 +50,8 @@ public class ShardSyncTaskIntegrationTest {
private static final String STREAM_NAME = "IntegrationTestStream02"; private static final String STREAM_NAME = "IntegrationTestStream02";
private static AmazonKinesis amazonKinesis; private static AmazonKinesis amazonKinesis;
private KinesisClientLeaseManager leaseManager; private LeaseRefresher leaseRefresher;
private LeaseManagerProxy leaseManagerProxy; private ShardDetector shardDetector;
/** /**
* @throws java.lang.Exception * @throws java.lang.Exception
@ -85,12 +87,13 @@ public class ShardSyncTaskIntegrationTest {
@Before @Before
public void setUp() throws Exception { public void setUp() throws Exception {
boolean useConsistentReads = true; boolean useConsistentReads = true;
leaseManager = leaseRefresher =
new KinesisClientDynamoDBLeaseManager("ShardSyncTaskIntegrationTest", new DynamoDBLeaseRefresher("ShardSyncTaskIntegrationTest",
AmazonDynamoDBClientBuilder.standard().withRegion(Regions.US_EAST_1).build(), AmazonDynamoDBClientBuilder.standard().withRegion(Regions.US_EAST_1).build(),
new DynamoDBLeaseSerializer(),
useConsistentReads); useConsistentReads);
leaseManagerProxy = new KinesisLeaseManagerProxy(amazonKinesis, STREAM_NAME, 500L, 50); shardDetector = new KinesisShardDetector(amazonKinesis, STREAM_NAME, 500L, 50);
} }
/** /**
@ -109,24 +112,24 @@ public class ShardSyncTaskIntegrationTest {
*/ */
@Test @Test
public final void testCall() throws DependencyException, InvalidStateException, ProvisionedThroughputException { public final void testCall() throws DependencyException, InvalidStateException, ProvisionedThroughputException {
if (!leaseManager.leaseTableExists()) { if (!leaseRefresher.leaseTableExists()) {
final Long readCapacity = 10L; final Long readCapacity = 10L;
final Long writeCapacity = 10L; final Long writeCapacity = 10L;
leaseManager.createLeaseTableIfNotExists(readCapacity, writeCapacity); leaseRefresher.createLeaseTableIfNotExists(readCapacity, writeCapacity);
} }
leaseManager.deleteAll(); leaseRefresher.deleteAll();
Set<String> shardIds = leaseManagerProxy.listShards().stream().map(Shard::getShardId).collect(Collectors.toSet()); Set<String> shardIds = shardDetector.listShards().stream().map(Shard::getShardId).collect(Collectors.toSet());
ShardSyncTask syncTask = new ShardSyncTask(leaseManagerProxy, ShardSyncTask syncTask = new ShardSyncTask(shardDetector,
leaseManager, leaseRefresher,
InitialPositionInStreamExtended.newInitialPosition(InitialPositionInStream.LATEST), InitialPositionInStreamExtended.newInitialPosition(InitialPositionInStream.LATEST),
false, false,
false, false,
0L); 0L);
syncTask.call(); syncTask.call();
List<KinesisClientLease> leases = leaseManager.listLeases(); List<Lease> leases = leaseRefresher.listLeases();
Set<String> leaseKeys = new HashSet<String>(); Set<String> leaseKeys = new HashSet<>();
for (KinesisClientLease lease : leases) { for (Lease lease : leases) {
leaseKeys.add(lease.getLeaseKey()); leaseKeys.add(lease.leaseKey());
} }
// Verify that all shardIds had leases for them // Verify that all shardIds had leases for them

View file

@ -50,7 +50,9 @@ import com.amazonaws.services.kinesis.model.SequenceNumberRange;
import com.amazonaws.services.kinesis.model.Shard; import com.amazonaws.services.kinesis.model.Shard;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.leases.ExceptionThrowingLeaseManager.ExceptionThrowingLeaseManagerMethods; import software.amazon.kinesis.leases.ExceptionThrowingLeaseRefresher.ExceptionThrowingLeaseRefresherMethods;
import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseRefresher;
import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseSerializer;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.LeasingException; import software.amazon.kinesis.leases.exceptions.LeasingException;
@ -70,30 +72,32 @@ public class ShardSyncerTest {
InitialPositionInStreamExtended.newInitialPosition(InitialPositionInStream.TRIM_HORIZON); InitialPositionInStreamExtended.newInitialPosition(InitialPositionInStream.TRIM_HORIZON);
private static final InitialPositionInStreamExtended INITIAL_POSITION_AT_TIMESTAMP = private static final InitialPositionInStreamExtended INITIAL_POSITION_AT_TIMESTAMP =
InitialPositionInStreamExtended.newInitialPositionAtTimestamp(new Date(1000L)); InitialPositionInStreamExtended.newInitialPositionAtTimestamp(new Date(1000L));
private static final int EXPONENT = 128;
private final boolean cleanupLeasesOfCompletedShards = true; private final boolean cleanupLeasesOfCompletedShards = true;
private AmazonDynamoDB ddbClient = DynamoDBEmbedded.create().amazonDynamoDB(); private AmazonDynamoDB ddbClient = DynamoDBEmbedded.create().amazonDynamoDB();
private DynamoDBLeaseManager<KinesisClientLease> dynamoDBLeaseManager = new KinesisClientDynamoDBLeaseManager("tempTestTable", ddbClient); private DynamoDBLeaseRefresher dynamoDBLeaseRefresher = new DynamoDBLeaseRefresher("tempTestTable", ddbClient,
private static final int EXPONENT = 128; new DynamoDBLeaseSerializer(), true);
/** /**
* Old/Obsolete max value of a sequence number (2^128 -1). * Old/Obsolete max value of a sequence number (2^128 -1).
*/ */
public static final BigInteger MAX_SEQUENCE_NUMBER = new BigInteger("2").pow(EXPONENT).subtract(BigInteger.ONE); public static final BigInteger MAX_SEQUENCE_NUMBER = new BigInteger("2").pow(EXPONENT).subtract(BigInteger.ONE);
@Mock @Mock
private LeaseManagerProxy leaseManagerProxy; private ShardDetector shardDetector;
@Before @Before
public void setUp() throws Exception { public void setUp() throws Exception {
boolean created = dynamoDBLeaseManager.createLeaseTableIfNotExists(1L, 1L); boolean created = dynamoDBLeaseRefresher.createLeaseTableIfNotExists(1L, 1L);
if (created) { if (created) {
log.info("New table created."); log.info("New table created.");
} }
dynamoDBLeaseManager.deleteAll(); dynamoDBLeaseRefresher.deleteAll();
} }
@After @After
public void tearDown() throws Exception { public void tearDown() throws Exception {
dynamoDBLeaseManager.deleteAll(); dynamoDBLeaseRefresher.deleteAll();
} }
/** /**
@ -102,7 +106,7 @@ public class ShardSyncerTest {
@Test @Test
public final void testDetermineNewLeasesToCreateNoShards() { public final void testDetermineNewLeasesToCreateNoShards() {
List<Shard> shards = new ArrayList<>(); List<Shard> shards = new ArrayList<>();
List<KinesisClientLease> leases = new ArrayList<>(); List<Lease> leases = new ArrayList<>();
assertTrue(ShardSyncer.determineNewLeasesToCreate(shards, leases, INITIAL_POSITION_LATEST).isEmpty()); assertTrue(ShardSyncer.determineNewLeasesToCreate(shards, leases, INITIAL_POSITION_LATEST).isEmpty());
} }
@ -113,7 +117,7 @@ public class ShardSyncerTest {
@Test @Test
public final void testDetermineNewLeasesToCreate0Leases0Reshards() { public final void testDetermineNewLeasesToCreate0Leases0Reshards() {
List<Shard> shards = new ArrayList<>(); List<Shard> shards = new ArrayList<>();
List<KinesisClientLease> currentLeases = new ArrayList<>(); List<Lease> currentLeases = new ArrayList<>();
SequenceNumberRange sequenceRange = ShardObjectHelper.newSequenceNumberRange("342980", null); SequenceNumberRange sequenceRange = ShardObjectHelper.newSequenceNumberRange("342980", null);
String shardId0 = "shardId-0"; String shardId0 = "shardId-0";
@ -122,14 +126,14 @@ public class ShardSyncerTest {
String shardId1 = "shardId-1"; String shardId1 = "shardId-1";
shards.add(ShardObjectHelper.newShard(shardId1, null, null, sequenceRange)); shards.add(ShardObjectHelper.newShard(shardId1, null, null, sequenceRange));
List<KinesisClientLease> newLeases = List<Lease> newLeases =
ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_LATEST); ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_LATEST);
assertEquals(2, newLeases.size()); assertEquals(2, newLeases.size());
Set<String> expectedLeaseShardIds = new HashSet<>(); Set<String> expectedLeaseShardIds = new HashSet<>();
expectedLeaseShardIds.add(shardId0); expectedLeaseShardIds.add(shardId0);
expectedLeaseShardIds.add(shardId1); expectedLeaseShardIds.add(shardId1);
for (KinesisClientLease lease : newLeases) { for (Lease lease : newLeases) {
assertTrue(expectedLeaseShardIds.contains(lease.getLeaseKey())); assertTrue(expectedLeaseShardIds.contains(lease.leaseKey()));
} }
} }
@ -140,7 +144,7 @@ public class ShardSyncerTest {
@Test @Test
public final void testDetermineNewLeasesToCreate0Leases0Reshards1Inconsistent() { public final void testDetermineNewLeasesToCreate0Leases0Reshards1Inconsistent() {
List<Shard> shards = new ArrayList<>(); List<Shard> shards = new ArrayList<>();
List<KinesisClientLease> currentLeases = new ArrayList<>(); List<Lease> currentLeases = new ArrayList<>();
SequenceNumberRange sequenceRange = ShardObjectHelper.newSequenceNumberRange("342980", null); SequenceNumberRange sequenceRange = ShardObjectHelper.newSequenceNumberRange("342980", null);
String shardId0 = "shardId-0"; String shardId0 = "shardId-0";
@ -155,14 +159,14 @@ public class ShardSyncerTest {
Set<String> inconsistentShardIds = new HashSet<>(); Set<String> inconsistentShardIds = new HashSet<>();
inconsistentShardIds.add(shardId2); inconsistentShardIds.add(shardId2);
List<KinesisClientLease> newLeases = List<Lease> newLeases =
ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_LATEST, inconsistentShardIds); ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_LATEST, inconsistentShardIds);
assertEquals(2, newLeases.size()); assertEquals(2, newLeases.size());
Set<String> expectedLeaseShardIds = new HashSet<>(); Set<String> expectedLeaseShardIds = new HashSet<>();
expectedLeaseShardIds.add(shardId0); expectedLeaseShardIds.add(shardId0);
expectedLeaseShardIds.add(shardId1); expectedLeaseShardIds.add(shardId1);
for (KinesisClientLease lease : newLeases) { for (Lease lease : newLeases) {
assertTrue(expectedLeaseShardIds.contains(lease.getLeaseKey())); assertTrue(expectedLeaseShardIds.contains(lease.leaseKey()));
} }
} }
@ -211,20 +215,20 @@ public class ShardSyncerTest {
IOException { IOException {
List<Shard> shards = constructShardListForGraphA(); List<Shard> shards = constructShardListForGraphA();
when(leaseManagerProxy.listShards()).thenReturn(shards); when(shardDetector.listShards()).thenReturn(shards);
ShardSyncer.checkAndCreateLeasesForNewShards(leaseManagerProxy, dynamoDBLeaseManager, INITIAL_POSITION_LATEST, ShardSyncer.checkAndCreateLeasesForNewShards(shardDetector, dynamoDBLeaseRefresher, INITIAL_POSITION_LATEST,
cleanupLeasesOfCompletedShards); cleanupLeasesOfCompletedShards);
List<KinesisClientLease> newLeases = dynamoDBLeaseManager.listLeases(); List<Lease> newLeases = dynamoDBLeaseRefresher.listLeases();
Set<String> expectedLeaseShardIds = new HashSet<>(); Set<String> expectedLeaseShardIds = new HashSet<>();
expectedLeaseShardIds.add("shardId-4"); expectedLeaseShardIds.add("shardId-4");
expectedLeaseShardIds.add("shardId-8"); expectedLeaseShardIds.add("shardId-8");
expectedLeaseShardIds.add("shardId-9"); expectedLeaseShardIds.add("shardId-9");
expectedLeaseShardIds.add("shardId-10"); expectedLeaseShardIds.add("shardId-10");
assertEquals(expectedLeaseShardIds.size(), newLeases.size()); assertEquals(expectedLeaseShardIds.size(), newLeases.size());
for (KinesisClientLease lease1 : newLeases) { for (Lease lease1 : newLeases) {
assertTrue(expectedLeaseShardIds.contains(lease1.getLeaseKey())); assertTrue(expectedLeaseShardIds.contains(lease1.leaseKey()));
assertEquals(ExtendedSequenceNumber.LATEST, lease1.getCheckpoint()); assertEquals(ExtendedSequenceNumber.LATEST, lease1.checkpoint());
} }
} }
@ -243,19 +247,19 @@ public class ShardSyncerTest {
IOException { IOException {
List<Shard> shards = constructShardListForGraphA(); List<Shard> shards = constructShardListForGraphA();
when(leaseManagerProxy.listShards()).thenReturn(shards); when(shardDetector.listShards()).thenReturn(shards);
ShardSyncer.checkAndCreateLeasesForNewShards(leaseManagerProxy, dynamoDBLeaseManager, INITIAL_POSITION_TRIM_HORIZON, ShardSyncer.checkAndCreateLeasesForNewShards(shardDetector, dynamoDBLeaseRefresher, INITIAL_POSITION_TRIM_HORIZON,
cleanupLeasesOfCompletedShards); cleanupLeasesOfCompletedShards);
List<KinesisClientLease> newLeases = dynamoDBLeaseManager.listLeases(); List<Lease> newLeases = dynamoDBLeaseRefresher.listLeases();
Set<String> expectedLeaseShardIds = new HashSet<>(); Set<String> expectedLeaseShardIds = new HashSet<>();
for (int i = 0; i < 11; i++) { for (int i = 0; i < 11; i++) {
expectedLeaseShardIds.add("shardId-" + i); expectedLeaseShardIds.add("shardId-" + i);
} }
assertEquals(expectedLeaseShardIds.size(), newLeases.size()); assertEquals(expectedLeaseShardIds.size(), newLeases.size());
for (KinesisClientLease lease1 : newLeases) { for (Lease lease1 : newLeases) {
assertTrue(expectedLeaseShardIds.contains(lease1.getLeaseKey())); assertTrue(expectedLeaseShardIds.contains(lease1.leaseKey()));
assertEquals(ExtendedSequenceNumber.TRIM_HORIZON, lease1.getCheckpoint()); assertEquals(ExtendedSequenceNumber.TRIM_HORIZON, lease1.checkpoint());
} }
} }
@ -272,19 +276,19 @@ public class ShardSyncerTest {
ProvisionedThroughputException, IOException { ProvisionedThroughputException, IOException {
List<Shard> shards = constructShardListForGraphA(); List<Shard> shards = constructShardListForGraphA();
when(leaseManagerProxy.listShards()).thenReturn(shards); when(shardDetector.listShards()).thenReturn(shards);
ShardSyncer.checkAndCreateLeasesForNewShards(leaseManagerProxy, dynamoDBLeaseManager, INITIAL_POSITION_AT_TIMESTAMP, ShardSyncer.checkAndCreateLeasesForNewShards(shardDetector, dynamoDBLeaseRefresher, INITIAL_POSITION_AT_TIMESTAMP,
cleanupLeasesOfCompletedShards); cleanupLeasesOfCompletedShards);
List<KinesisClientLease> newLeases = dynamoDBLeaseManager.listLeases(); List<Lease> newLeases = dynamoDBLeaseRefresher.listLeases();
Set<String> expectedLeaseShardIds = new HashSet<>(); Set<String> expectedLeaseShardIds = new HashSet<>();
for (int i = 0; i < 11; i++) { for (int i = 0; i < 11; i++) {
expectedLeaseShardIds.add("shardId-" + i); expectedLeaseShardIds.add("shardId-" + i);
} }
assertEquals(expectedLeaseShardIds.size(), newLeases.size()); assertEquals(expectedLeaseShardIds.size(), newLeases.size());
for (KinesisClientLease lease1 : newLeases) { for (Lease lease1 : newLeases) {
assertTrue(expectedLeaseShardIds.contains(lease1.getLeaseKey())); assertTrue(expectedLeaseShardIds.contains(lease1.leaseKey()));
assertEquals(ExtendedSequenceNumber.AT_TIMESTAMP, lease1.getCheckpoint()); assertEquals(ExtendedSequenceNumber.AT_TIMESTAMP, lease1.checkpoint());
} }
} }
@ -304,9 +308,9 @@ public class ShardSyncerTest {
range.setEndingSequenceNumber(null); range.setEndingSequenceNumber(null);
shards.get(3).setSequenceNumberRange(range); shards.get(3).setSequenceNumberRange(range);
when(leaseManagerProxy.listShards()).thenReturn(shards); when(shardDetector.listShards()).thenReturn(shards);
ShardSyncer.checkAndCreateLeasesForNewShards(leaseManagerProxy, dynamoDBLeaseManager, INITIAL_POSITION_TRIM_HORIZON, ShardSyncer.checkAndCreateLeasesForNewShards(shardDetector, dynamoDBLeaseRefresher, INITIAL_POSITION_TRIM_HORIZON,
cleanupLeasesOfCompletedShards); cleanupLeasesOfCompletedShards);
} }
@ -328,19 +332,19 @@ public class ShardSyncerTest {
range.setEndingSequenceNumber(null); range.setEndingSequenceNumber(null);
shard.setSequenceNumberRange(range); shard.setSequenceNumberRange(range);
when(leaseManagerProxy.listShards()).thenReturn(shards); when(shardDetector.listShards()).thenReturn(shards);
ShardSyncer.checkAndCreateLeasesForNewShards(leaseManagerProxy, dynamoDBLeaseManager, INITIAL_POSITION_LATEST, ShardSyncer.checkAndCreateLeasesForNewShards(shardDetector, dynamoDBLeaseRefresher, INITIAL_POSITION_LATEST,
cleanupLeasesOfCompletedShards, true); cleanupLeasesOfCompletedShards, true);
List<KinesisClientLease> newLeases = dynamoDBLeaseManager.listLeases(); List<Lease> newLeases = dynamoDBLeaseRefresher.listLeases();
Set<String> expectedLeaseShardIds = new HashSet<>(); Set<String> expectedLeaseShardIds = new HashSet<>();
expectedLeaseShardIds.add("shardId-4"); expectedLeaseShardIds.add("shardId-4");
expectedLeaseShardIds.add("shardId-5"); expectedLeaseShardIds.add("shardId-5");
expectedLeaseShardIds.add("shardId-8"); expectedLeaseShardIds.add("shardId-8");
assertEquals(expectedLeaseShardIds.size(), newLeases.size()); assertEquals(expectedLeaseShardIds.size(), newLeases.size());
for (KinesisClientLease lease1 : newLeases) { for (Lease lease1 : newLeases) {
assertTrue(expectedLeaseShardIds.contains(lease1.getLeaseKey())); assertTrue(expectedLeaseShardIds.contains(lease1.leaseKey()));
assertEquals(ExtendedSequenceNumber.LATEST, lease1.getCheckpoint()); assertEquals(ExtendedSequenceNumber.LATEST, lease1.checkpoint());
} }
} }
@ -372,14 +376,14 @@ public class ShardSyncerTest {
public final void testCheckAndCreateLeasesForNewShardsAtTrimHorizonAndClosedShardWithDeleteLeaseExceptions() public final void testCheckAndCreateLeasesForNewShardsAtTrimHorizonAndClosedShardWithDeleteLeaseExceptions()
throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException, throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException,
IOException { IOException {
// Define the max calling count for lease manager methods. // Define the max calling count for lease refresher methods.
// From the Shard Graph, the max count of calling could be 10 // From the Shard Graph, the max count of calling could be 10
int maxCallingCount = 10; int maxCallingCount = 10;
for (int c = 1; c <= maxCallingCount; c = c + 2) { for (int c = 1; c <= maxCallingCount; c = c + 2) {
testCheckAndCreateLeasesForNewShardsAtSpecifiedPositionAndClosedShardImpl( testCheckAndCreateLeasesForNewShardsAtSpecifiedPositionAndClosedShardImpl(
ExceptionThrowingLeaseManagerMethods.DELETELEASE, c, INITIAL_POSITION_TRIM_HORIZON); ExceptionThrowingLeaseRefresherMethods.DELETELEASE, c, INITIAL_POSITION_TRIM_HORIZON);
// Need to clean up lease manager every time after calling ShardSyncer // Need to clean up lease refresher every time after calling ShardSyncer
dynamoDBLeaseManager.deleteAll(); dynamoDBLeaseRefresher.deleteAll();
} }
} }
@ -396,14 +400,14 @@ public class ShardSyncerTest {
public final void testCheckAndCreateLeasesForNewShardsAtTrimHorizonAndClosedShardWithListLeasesExceptions() public final void testCheckAndCreateLeasesForNewShardsAtTrimHorizonAndClosedShardWithListLeasesExceptions()
throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException, throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException,
IOException { IOException {
// Define the max calling count for lease manager methods. // Define the max calling count for lease refresher methods.
// From the Shard Graph, the max count of calling could be 10 // From the Shard Graph, the max count of calling could be 10
int maxCallingCount = 10; int maxCallingCount = 10;
for (int c = 1; c <= maxCallingCount; c = c + 2) { for (int c = 1; c <= maxCallingCount; c = c + 2) {
testCheckAndCreateLeasesForNewShardsAtSpecifiedPositionAndClosedShardImpl( testCheckAndCreateLeasesForNewShardsAtSpecifiedPositionAndClosedShardImpl(
ExceptionThrowingLeaseManagerMethods.LISTLEASES, c, INITIAL_POSITION_TRIM_HORIZON); ExceptionThrowingLeaseRefresherMethods.LISTLEASES, c, INITIAL_POSITION_TRIM_HORIZON);
// Need to clean up lease manager every time after calling ShardSyncer // Need to clean up lease refresher every time after calling ShardSyncer
dynamoDBLeaseManager.deleteAll(); dynamoDBLeaseRefresher.deleteAll();
} }
} }
@ -420,34 +424,34 @@ public class ShardSyncerTest {
public final void testCheckAndCreateLeasesForNewShardsAtTrimHorizonAndClosedShardWithCreateLeaseExceptions() public final void testCheckAndCreateLeasesForNewShardsAtTrimHorizonAndClosedShardWithCreateLeaseExceptions()
throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException, throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException,
IOException { IOException {
// Define the max calling count for lease manager methods. // Define the max calling count for lease refresher methods.
// From the Shard Graph, the max count of calling could be 10 // From the Shard Graph, the max count of calling could be 10
int maxCallingCount = 5; int maxCallingCount = 5;
for (int c = 1; c <= maxCallingCount; c = c + 2) { for (int c = 1; c <= maxCallingCount; c = c + 2) {
testCheckAndCreateLeasesForNewShardsAtSpecifiedPositionAndClosedShardImpl( testCheckAndCreateLeasesForNewShardsAtSpecifiedPositionAndClosedShardImpl(
ExceptionThrowingLeaseManagerMethods.CREATELEASEIFNOTEXISTS, c,INITIAL_POSITION_TRIM_HORIZON); ExceptionThrowingLeaseRefresherMethods.CREATELEASEIFNOTEXISTS, c,INITIAL_POSITION_TRIM_HORIZON);
// Need to clean up lease manager every time after calling ShardSyncer // Need to clean up lease refresher every time after calling ShardSyncer
dynamoDBLeaseManager.deleteAll(); dynamoDBLeaseRefresher.deleteAll();
} }
} }
// Try catch leaseException for different lease manager methods and eventually let it succeed. // Try catch leaseException for different lease refresher methods and eventually let it succeed.
// This would not throw any exceptions if: // This would not throw any exceptions if:
// 1). exceptionMethod equals to null or NONE. // 1). exceptionMethod equals to null or NONE.
// 2). exceptionTime is a very big or negative value. // 2). exceptionTime is a very big or negative value.
private void retryCheckAndCreateLeaseForNewShards(ExceptionThrowingLeaseManagerMethods exceptionMethod, private void retryCheckAndCreateLeaseForNewShards(ExceptionThrowingLeaseRefresherMethods exceptionMethod,
int exceptionTime, InitialPositionInStreamExtended position) int exceptionTime, InitialPositionInStreamExtended position)
throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException { throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException {
if (exceptionMethod != null) { if (exceptionMethod != null) {
ExceptionThrowingLeaseManager exceptionThrowingLeaseManager = ExceptionThrowingLeaseRefresher exceptionThrowingLeaseRefresher =
new ExceptionThrowingLeaseManager(dynamoDBLeaseManager); new ExceptionThrowingLeaseRefresher(dynamoDBLeaseRefresher);
// Set exception and throwing time for exceptionThrowingManager. // Set exception and throwing time for exceptionThrowingRefresher.
exceptionThrowingLeaseManager.setLeaseLeaseManagerThrowingExceptionScenario(exceptionMethod, exceptionTime); exceptionThrowingLeaseRefresher.leaseRefresherThrowingExceptionScenario(exceptionMethod, exceptionTime);
// Only need to try two times. // Only need to try two times.
for (int i = 1; i <= 2; i++) { for (int i = 1; i <= 2; i++) {
try { try {
ShardSyncer.checkAndCreateLeasesForNewShards(leaseManagerProxy, ShardSyncer.checkAndCreateLeasesForNewShards(shardDetector,
exceptionThrowingLeaseManager, exceptionThrowingLeaseRefresher,
position, position,
cleanupLeasesOfCompletedShards); cleanupLeasesOfCompletedShards);
return; return;
@ -455,11 +459,11 @@ public class ShardSyncerTest {
log.debug("Catch leasing exception", e); log.debug("Catch leasing exception", e);
} }
// Clear throwing exception scenario every time after calling ShardSyncer // Clear throwing exception scenario every time after calling ShardSyncer
exceptionThrowingLeaseManager.clearLeaseManagerThrowingExceptionScenario(); exceptionThrowingLeaseRefresher.clearLeaseRefresherThrowingExceptionScenario();
} }
} else { } else {
ShardSyncer.checkAndCreateLeasesForNewShards(leaseManagerProxy, ShardSyncer.checkAndCreateLeasesForNewShards(shardDetector,
dynamoDBLeaseManager, dynamoDBLeaseRefresher,
position, position,
cleanupLeasesOfCompletedShards); cleanupLeasesOfCompletedShards);
} }
@ -493,15 +497,15 @@ public class ShardSyncerTest {
public final void testCheckAndCreateLeasesForNewShardsAtTimestampAndClosedShardWithDeleteLeaseExceptions() public final void testCheckAndCreateLeasesForNewShardsAtTimestampAndClosedShardWithDeleteLeaseExceptions()
throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException, throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException,
IOException { IOException {
// Define the max calling count for lease manager methods. // Define the max calling count for lease refresher methods.
// From the Shard Graph, the max count of calling could be 10 // From the Shard Graph, the max count of calling could be 10
int maxCallingCount = 10; int maxCallingCount = 10;
for (int c = 1; c <= maxCallingCount; c = c + 2) { for (int c = 1; c <= maxCallingCount; c = c + 2) {
testCheckAndCreateLeasesForNewShardsAtSpecifiedPositionAndClosedShardImpl( testCheckAndCreateLeasesForNewShardsAtSpecifiedPositionAndClosedShardImpl(
ExceptionThrowingLeaseManagerMethods.DELETELEASE, ExceptionThrowingLeaseRefresherMethods.DELETELEASE,
c, INITIAL_POSITION_AT_TIMESTAMP); c, INITIAL_POSITION_AT_TIMESTAMP);
// Need to clean up lease manager every time after calling ShardSyncer // Need to clean up lease refresher every time after calling ShardSyncer
dynamoDBLeaseManager.deleteAll(); dynamoDBLeaseRefresher.deleteAll();
} }
} }
@ -516,15 +520,15 @@ public class ShardSyncerTest {
public final void testCheckAndCreateLeasesForNewShardsAtTimestampAndClosedShardWithListLeasesExceptions() public final void testCheckAndCreateLeasesForNewShardsAtTimestampAndClosedShardWithListLeasesExceptions()
throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException, throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException,
IOException { IOException {
// Define the max calling count for lease manager methods. // Define the max calling count for lease refresher methods.
// From the Shard Graph, the max count of calling could be 10 // From the Shard Graph, the max count of calling could be 10
int maxCallingCount = 10; int maxCallingCount = 10;
for (int c = 1; c <= maxCallingCount; c = c + 2) { for (int c = 1; c <= maxCallingCount; c = c + 2) {
testCheckAndCreateLeasesForNewShardsAtSpecifiedPositionAndClosedShardImpl( testCheckAndCreateLeasesForNewShardsAtSpecifiedPositionAndClosedShardImpl(
ExceptionThrowingLeaseManagerMethods.LISTLEASES, ExceptionThrowingLeaseRefresherMethods.LISTLEASES,
c, INITIAL_POSITION_AT_TIMESTAMP); c, INITIAL_POSITION_AT_TIMESTAMP);
// Need to clean up lease manager every time after calling ShardSyncer // Need to clean up lease refresher every time after calling ShardSyncer
dynamoDBLeaseManager.deleteAll(); dynamoDBLeaseRefresher.deleteAll();
} }
} }
@ -539,21 +543,21 @@ public class ShardSyncerTest {
public final void testCheckAndCreateLeasesForNewShardsAtTimestampAndClosedShardWithCreateLeaseExceptions() public final void testCheckAndCreateLeasesForNewShardsAtTimestampAndClosedShardWithCreateLeaseExceptions()
throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException, throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException,
IOException { IOException {
// Define the max calling count for lease manager methods. // Define the max calling count for lease refresher methods.
// From the Shard Graph, the max count of calling could be 10 // From the Shard Graph, the max count of calling could be 10
int maxCallingCount = 5; int maxCallingCount = 5;
for (int c = 1; c <= maxCallingCount; c = c + 2) { for (int c = 1; c <= maxCallingCount; c = c + 2) {
testCheckAndCreateLeasesForNewShardsAtSpecifiedPositionAndClosedShardImpl( testCheckAndCreateLeasesForNewShardsAtSpecifiedPositionAndClosedShardImpl(
ExceptionThrowingLeaseManagerMethods.CREATELEASEIFNOTEXISTS, ExceptionThrowingLeaseRefresherMethods.CREATELEASEIFNOTEXISTS,
c, INITIAL_POSITION_AT_TIMESTAMP); c, INITIAL_POSITION_AT_TIMESTAMP);
// Need to clean up lease manager every time after calling ShardSyncer // Need to clean up lease refresher every time after calling ShardSyncer
dynamoDBLeaseManager.deleteAll(); dynamoDBLeaseRefresher.deleteAll();
} }
} }
// Real implementation of testing CheckAndCreateLeasesForNewShards with different leaseManager types. // Real implementation of testing CheckAndCreateLeasesForNewShards with different leaseRefresher types.
private void testCheckAndCreateLeasesForNewShardsAtSpecifiedPositionAndClosedShardImpl( private void testCheckAndCreateLeasesForNewShardsAtSpecifiedPositionAndClosedShardImpl(
ExceptionThrowingLeaseManagerMethods exceptionMethod, ExceptionThrowingLeaseRefresherMethods exceptionMethod,
int exceptionTime, int exceptionTime,
InitialPositionInStreamExtended position) InitialPositionInStreamExtended position)
throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException, throws KinesisClientLibIOException, DependencyException, InvalidStateException, ProvisionedThroughputException,
@ -562,39 +566,39 @@ public class ShardSyncerTest {
new ExtendedSequenceNumber(position.getInitialPositionInStream().toString()); new ExtendedSequenceNumber(position.getInitialPositionInStream().toString());
List<Shard> shards = constructShardListForGraphA(); List<Shard> shards = constructShardListForGraphA();
when(leaseManagerProxy.listShards()).thenReturn(shards); when(shardDetector.listShards()).thenReturn(shards);
retryCheckAndCreateLeaseForNewShards(exceptionMethod, exceptionTime, position); retryCheckAndCreateLeaseForNewShards(exceptionMethod, exceptionTime, position);
List<KinesisClientLease> newLeases = dynamoDBLeaseManager.listLeases(); List<Lease> newLeases = dynamoDBLeaseRefresher.listLeases();
Map<String, ExtendedSequenceNumber> expectedShardIdToCheckpointMap = new HashMap<>(); Map<String, ExtendedSequenceNumber> expectedShardIdToCheckpointMap = new HashMap<>();
for (int i = 0; i < 11; i++) { for (int i = 0; i < 11; i++) {
expectedShardIdToCheckpointMap.put("shardId-" + i, extendedSequenceNumber); expectedShardIdToCheckpointMap.put("shardId-" + i, extendedSequenceNumber);
} }
assertEquals(expectedShardIdToCheckpointMap.size(), newLeases.size()); assertEquals(expectedShardIdToCheckpointMap.size(), newLeases.size());
for (KinesisClientLease lease1 : newLeases) { for (Lease lease1 : newLeases) {
ExtendedSequenceNumber expectedCheckpoint = expectedShardIdToCheckpointMap.get(lease1.getLeaseKey()); ExtendedSequenceNumber expectedCheckpoint = expectedShardIdToCheckpointMap.get(lease1.leaseKey());
assertNotNull(expectedCheckpoint); assertNotNull(expectedCheckpoint);
assertEquals(expectedCheckpoint, lease1.getCheckpoint()); assertEquals(expectedCheckpoint, lease1.checkpoint());
} }
KinesisClientLease closedShardLease = dynamoDBLeaseManager.getLease("shardId-0"); Lease closedShardLease = dynamoDBLeaseRefresher.getLease("shardId-0");
closedShardLease.setCheckpoint(ExtendedSequenceNumber.SHARD_END); closedShardLease.checkpoint(ExtendedSequenceNumber.SHARD_END);
dynamoDBLeaseManager.updateLease(closedShardLease); dynamoDBLeaseRefresher.updateLease(closedShardLease);
expectedShardIdToCheckpointMap.remove(closedShardLease.getLeaseKey()); expectedShardIdToCheckpointMap.remove(closedShardLease.leaseKey());
KinesisClientLease childShardLease = dynamoDBLeaseManager.getLease("shardId-6"); Lease childShardLease = dynamoDBLeaseRefresher.getLease("shardId-6");
childShardLease.setCheckpoint(new ExtendedSequenceNumber("34290")); childShardLease.checkpoint(new ExtendedSequenceNumber("34290"));
dynamoDBLeaseManager.updateLease(childShardLease); dynamoDBLeaseRefresher.updateLease(childShardLease);
expectedShardIdToCheckpointMap.put(childShardLease.getLeaseKey(), new ExtendedSequenceNumber("34290")); expectedShardIdToCheckpointMap.put(childShardLease.leaseKey(), new ExtendedSequenceNumber("34290"));
retryCheckAndCreateLeaseForNewShards(exceptionMethod, exceptionTime, position); retryCheckAndCreateLeaseForNewShards(exceptionMethod, exceptionTime, position);
newLeases = dynamoDBLeaseManager.listLeases(); newLeases = dynamoDBLeaseRefresher.listLeases();
assertEquals(expectedShardIdToCheckpointMap.size(), newLeases.size()); assertEquals(expectedShardIdToCheckpointMap.size(), newLeases.size());
for (KinesisClientLease lease1 : newLeases) { for (Lease lease1 : newLeases) {
ExtendedSequenceNumber expectedCheckpoint = expectedShardIdToCheckpointMap.get(lease1.getLeaseKey()); ExtendedSequenceNumber expectedCheckpoint = expectedShardIdToCheckpointMap.get(lease1.leaseKey());
assertNotNull(expectedCheckpoint); assertNotNull(expectedCheckpoint);
assertEquals(expectedCheckpoint, lease1.getCheckpoint()); assertEquals(expectedCheckpoint, lease1.checkpoint());
} }
} }
@ -612,15 +616,15 @@ public class ShardSyncerTest {
throws DependencyException, InvalidStateException, ProvisionedThroughputException, IOException, throws DependencyException, InvalidStateException, ProvisionedThroughputException, IOException,
KinesisClientLibIOException { KinesisClientLibIOException {
String garbageShardId = "shardId-garbage-001"; String garbageShardId = "shardId-garbage-001";
KinesisClientLease garbageLease = ShardSyncer.newKCLLease(ShardObjectHelper.newShard(garbageShardId, Lease garbageLease = ShardSyncer.newKCLLease(ShardObjectHelper.newShard(garbageShardId,
null, null,
null, null,
ShardObjectHelper.newSequenceNumberRange("101", null))); ShardObjectHelper.newSequenceNumberRange("101", null)));
garbageLease.setCheckpoint(new ExtendedSequenceNumber("999")); garbageLease.checkpoint(new ExtendedSequenceNumber("999"));
dynamoDBLeaseManager.createLeaseIfNotExists(garbageLease); dynamoDBLeaseRefresher.createLeaseIfNotExists(garbageLease);
assertEquals(garbageShardId, dynamoDBLeaseManager.getLease(garbageShardId).getLeaseKey()); assertEquals(garbageShardId, dynamoDBLeaseRefresher.getLease(garbageShardId).leaseKey());
testBootstrapShardLeasesAtStartingPosition(INITIAL_POSITION_LATEST); testBootstrapShardLeasesAtStartingPosition(INITIAL_POSITION_LATEST);
assertNull(dynamoDBLeaseManager.getLease(garbageShardId)); assertNull(dynamoDBLeaseRefresher.getLease(garbageShardId));
} }
private void testBootstrapShardLeasesAtStartingPosition(InitialPositionInStreamExtended initialPosition) private void testBootstrapShardLeasesAtStartingPosition(InitialPositionInStreamExtended initialPosition)
@ -634,19 +638,19 @@ public class ShardSyncerTest {
String shardId1 = "shardId-1"; String shardId1 = "shardId-1";
shards.add(ShardObjectHelper.newShard(shardId1, null, null, sequenceRange)); shards.add(ShardObjectHelper.newShard(shardId1, null, null, sequenceRange));
when(leaseManagerProxy.listShards()).thenReturn(shards); when(shardDetector.listShards()).thenReturn(shards);
ShardSyncer.bootstrapShardLeases(leaseManagerProxy, dynamoDBLeaseManager, initialPosition, cleanupLeasesOfCompletedShards, ShardSyncer.bootstrapShardLeases(shardDetector, dynamoDBLeaseRefresher, initialPosition, cleanupLeasesOfCompletedShards,
false); false);
List<KinesisClientLease> newLeases = dynamoDBLeaseManager.listLeases(); List<Lease> newLeases = dynamoDBLeaseRefresher.listLeases();
assertEquals(2, newLeases.size()); assertEquals(2, newLeases.size());
Set<String> expectedLeaseShardIds = new HashSet<String>(); Set<String> expectedLeaseShardIds = new HashSet<String>();
expectedLeaseShardIds.add(shardId0); expectedLeaseShardIds.add(shardId0);
expectedLeaseShardIds.add(shardId1); expectedLeaseShardIds.add(shardId1);
for (KinesisClientLease lease1 : newLeases) { for (Lease lease1 : newLeases) {
assertTrue(expectedLeaseShardIds.contains(lease1.getLeaseKey())); assertTrue(expectedLeaseShardIds.contains(lease1.leaseKey()));
assertEquals(new ExtendedSequenceNumber(initialPosition.getInitialPositionInStream().toString()), assertEquals(new ExtendedSequenceNumber(initialPosition.getInitialPositionInStream().toString()),
lease1.getCheckpoint()); lease1.checkpoint());
} }
} }
@ -656,7 +660,7 @@ public class ShardSyncerTest {
@Test @Test
public final void testDetermineNewLeasesToCreateStartingPosition() { public final void testDetermineNewLeasesToCreateStartingPosition() {
List<Shard> shards = new ArrayList<>(); List<Shard> shards = new ArrayList<>();
List<KinesisClientLease> currentLeases = new ArrayList<>(); List<Lease> currentLeases = new ArrayList<>();
SequenceNumberRange sequenceRange = ShardObjectHelper.newSequenceNumberRange("342980", null); SequenceNumberRange sequenceRange = ShardObjectHelper.newSequenceNumberRange("342980", null);
String shardId0 = "shardId-0"; String shardId0 = "shardId-0";
@ -670,16 +674,16 @@ public class ShardSyncerTest {
initialPositions.add(INITIAL_POSITION_TRIM_HORIZON); initialPositions.add(INITIAL_POSITION_TRIM_HORIZON);
for (InitialPositionInStreamExtended initialPosition : initialPositions) { for (InitialPositionInStreamExtended initialPosition : initialPositions) {
List<KinesisClientLease> newLeases = List<Lease> newLeases =
ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, initialPosition); ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, initialPosition);
assertEquals(2, newLeases.size()); assertEquals(2, newLeases.size());
Set<String> expectedLeaseShardIds = new HashSet<String>(); Set<String> expectedLeaseShardIds = new HashSet<String>();
expectedLeaseShardIds.add(shardId0); expectedLeaseShardIds.add(shardId0);
expectedLeaseShardIds.add(shardId1); expectedLeaseShardIds.add(shardId1);
for (KinesisClientLease lease : newLeases) { for (Lease lease : newLeases) {
assertTrue(expectedLeaseShardIds.contains(lease.getLeaseKey())); assertTrue(expectedLeaseShardIds.contains(lease.leaseKey()));
assertEquals(new ExtendedSequenceNumber(initialPosition.getInitialPositionInStream().toString()), assertEquals(new ExtendedSequenceNumber(initialPosition.getInitialPositionInStream().toString()),
lease.getCheckpoint()); lease.checkpoint());
} }
} }
} }
@ -690,7 +694,7 @@ public class ShardSyncerTest {
@Test @Test
public final void testDetermineNewLeasesToCreateIgnoreClosedShard() { public final void testDetermineNewLeasesToCreateIgnoreClosedShard() {
List<Shard> shards = new ArrayList<>(); List<Shard> shards = new ArrayList<>();
List<KinesisClientLease> currentLeases = new ArrayList<>(); List<Lease> currentLeases = new ArrayList<>();
shards.add(ShardObjectHelper.newShard("shardId-0", shards.add(ShardObjectHelper.newShard("shardId-0",
null, null,
@ -702,10 +706,10 @@ public class ShardSyncerTest {
null, null,
ShardObjectHelper.newSequenceNumberRange("405", null))); ShardObjectHelper.newSequenceNumberRange("405", null)));
List<KinesisClientLease> newLeases = List<Lease> newLeases =
ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_LATEST); ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_LATEST);
assertEquals(1, newLeases.size()); assertEquals(1, newLeases.size());
assertEquals(lastShardId, newLeases.get(0).getLeaseKey()); assertEquals(lastShardId, newLeases.get(0).leaseKey());
} }
/** /**
@ -721,13 +725,13 @@ public class ShardSyncerTest {
@Test @Test
public final void testDetermineNewLeasesToCreateSplitMergeLatest1() { public final void testDetermineNewLeasesToCreateSplitMergeLatest1() {
List<Shard> shards = constructShardListForGraphA(); List<Shard> shards = constructShardListForGraphA();
List<KinesisClientLease> currentLeases = new ArrayList<>(); List<Lease> currentLeases = new ArrayList<>();
currentLeases.add(newLease("shardId-3")); currentLeases.add(newLease("shardId-3"));
currentLeases.add(newLease("shardId-4")); currentLeases.add(newLease("shardId-4"));
currentLeases.add(newLease("shardId-5")); currentLeases.add(newLease("shardId-5"));
List<KinesisClientLease> newLeases = List<Lease> newLeases =
ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_LATEST); ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_LATEST);
Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = new HashMap<>(); Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = new HashMap<>();
expectedShardIdCheckpointMap.put("shardId-8", ExtendedSequenceNumber.TRIM_HORIZON); expectedShardIdCheckpointMap.put("shardId-8", ExtendedSequenceNumber.TRIM_HORIZON);
@ -738,10 +742,10 @@ public class ShardSyncerTest {
expectedShardIdCheckpointMap.put("shardId-7", ExtendedSequenceNumber.TRIM_HORIZON); expectedShardIdCheckpointMap.put("shardId-7", ExtendedSequenceNumber.TRIM_HORIZON);
assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size()); assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size());
for (KinesisClientLease lease : newLeases) { for (Lease lease : newLeases) {
assertTrue("Unexpected lease: " + lease, assertTrue("Unexpected lease: " + lease,
expectedShardIdCheckpointMap.containsKey(lease.getLeaseKey())); expectedShardIdCheckpointMap.containsKey(lease.leaseKey()));
assertEquals(expectedShardIdCheckpointMap.get(lease.getLeaseKey()), lease.getCheckpoint()); assertEquals(expectedShardIdCheckpointMap.get(lease.leaseKey()), lease.checkpoint());
} }
} }
@ -758,13 +762,13 @@ public class ShardSyncerTest {
@Test @Test
public final void testDetermineNewLeasesToCreateSplitMergeLatest2() { public final void testDetermineNewLeasesToCreateSplitMergeLatest2() {
List<Shard> shards = constructShardListForGraphA(); List<Shard> shards = constructShardListForGraphA();
List<KinesisClientLease> currentLeases = new ArrayList<>(); List<Lease> currentLeases = new ArrayList<>();
currentLeases.add(newLease("shardId-4")); currentLeases.add(newLease("shardId-4"));
currentLeases.add(newLease("shardId-5")); currentLeases.add(newLease("shardId-5"));
currentLeases.add(newLease("shardId-7")); currentLeases.add(newLease("shardId-7"));
List<KinesisClientLease> newLeases = List<Lease> newLeases =
ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_LATEST); ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_LATEST);
Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = new HashMap<>(); Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = new HashMap<>();
expectedShardIdCheckpointMap.put("shardId-8", ExtendedSequenceNumber.TRIM_HORIZON); expectedShardIdCheckpointMap.put("shardId-8", ExtendedSequenceNumber.TRIM_HORIZON);
@ -773,10 +777,10 @@ public class ShardSyncerTest {
expectedShardIdCheckpointMap.put("shardId-6", ExtendedSequenceNumber.LATEST); expectedShardIdCheckpointMap.put("shardId-6", ExtendedSequenceNumber.LATEST);
assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size()); assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size());
for (KinesisClientLease lease : newLeases) { for (Lease lease : newLeases) {
assertTrue("Unexpected lease: " + lease, assertTrue("Unexpected lease: " + lease,
expectedShardIdCheckpointMap.containsKey(lease.getLeaseKey())); expectedShardIdCheckpointMap.containsKey(lease.leaseKey()));
assertEquals(expectedShardIdCheckpointMap.get(lease.getLeaseKey()), lease.getCheckpoint()); assertEquals(expectedShardIdCheckpointMap.get(lease.leaseKey()), lease.checkpoint());
} }
} }
@ -793,13 +797,13 @@ public class ShardSyncerTest {
@Test @Test
public final void testDetermineNewLeasesToCreateSplitMergeHorizon1() { public final void testDetermineNewLeasesToCreateSplitMergeHorizon1() {
List<Shard> shards = constructShardListForGraphA(); List<Shard> shards = constructShardListForGraphA();
List<KinesisClientLease> currentLeases = new ArrayList<>(); List<Lease> currentLeases = new ArrayList<>();
currentLeases.add(newLease("shardId-3")); currentLeases.add(newLease("shardId-3"));
currentLeases.add(newLease("shardId-4")); currentLeases.add(newLease("shardId-4"));
currentLeases.add(newLease("shardId-5")); currentLeases.add(newLease("shardId-5"));
List<KinesisClientLease> newLeases = List<Lease> newLeases =
ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_TRIM_HORIZON); ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_TRIM_HORIZON);
Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = new HashMap<>(); Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = new HashMap<>();
expectedShardIdCheckpointMap.put("shardId-8", ExtendedSequenceNumber.TRIM_HORIZON); expectedShardIdCheckpointMap.put("shardId-8", ExtendedSequenceNumber.TRIM_HORIZON);
@ -812,10 +816,10 @@ public class ShardSyncerTest {
expectedShardIdCheckpointMap.put("shardId-1", ExtendedSequenceNumber.TRIM_HORIZON); expectedShardIdCheckpointMap.put("shardId-1", ExtendedSequenceNumber.TRIM_HORIZON);
assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size()); assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size());
for (KinesisClientLease lease : newLeases) { for (Lease lease : newLeases) {
assertTrue("Unexpected lease: " + lease, assertTrue("Unexpected lease: " + lease,
expectedShardIdCheckpointMap.containsKey(lease.getLeaseKey())); expectedShardIdCheckpointMap.containsKey(lease.leaseKey()));
assertEquals(expectedShardIdCheckpointMap.get(lease.getLeaseKey()), lease.getCheckpoint()); assertEquals(expectedShardIdCheckpointMap.get(lease.leaseKey()), lease.checkpoint());
} }
} }
@ -832,13 +836,13 @@ public class ShardSyncerTest {
@Test @Test
public final void testDetermineNewLeasesToCreateSplitMergeHorizon2() { public final void testDetermineNewLeasesToCreateSplitMergeHorizon2() {
List<Shard> shards = constructShardListForGraphA(); List<Shard> shards = constructShardListForGraphA();
List<KinesisClientLease> currentLeases = new ArrayList<>(); List<Lease> currentLeases = new ArrayList<>();
currentLeases.add(newLease("shardId-4")); currentLeases.add(newLease("shardId-4"));
currentLeases.add(newLease("shardId-5")); currentLeases.add(newLease("shardId-5"));
currentLeases.add(newLease("shardId-7")); currentLeases.add(newLease("shardId-7"));
List<KinesisClientLease> newLeases = List<Lease> newLeases =
ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_TRIM_HORIZON); ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_TRIM_HORIZON);
Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap =
new HashMap<>(); new HashMap<>();
@ -850,10 +854,10 @@ public class ShardSyncerTest {
expectedShardIdCheckpointMap.put("shardId-1", ExtendedSequenceNumber.TRIM_HORIZON); expectedShardIdCheckpointMap.put("shardId-1", ExtendedSequenceNumber.TRIM_HORIZON);
assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size()); assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size());
for (KinesisClientLease lease : newLeases) { for (Lease lease : newLeases) {
assertTrue("Unexpected lease: " + lease, assertTrue("Unexpected lease: " + lease,
expectedShardIdCheckpointMap.containsKey(lease.getLeaseKey())); expectedShardIdCheckpointMap.containsKey(lease.leaseKey()));
assertEquals(expectedShardIdCheckpointMap.get(lease.getLeaseKey()), lease.getCheckpoint()); assertEquals(expectedShardIdCheckpointMap.get(lease.leaseKey()), lease.checkpoint());
} }
} }
@ -866,8 +870,8 @@ public class ShardSyncerTest {
@Test @Test
public final void testDetermineNewLeasesToCreateGraphBNoInitialLeasesTrim() { public final void testDetermineNewLeasesToCreateGraphBNoInitialLeasesTrim() {
List<Shard> shards = constructShardListForGraphB(); List<Shard> shards = constructShardListForGraphB();
List<KinesisClientLease> currentLeases = new ArrayList<>(); List<Lease> currentLeases = new ArrayList<>();
List<KinesisClientLease> newLeases = List<Lease> newLeases =
ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_TRIM_HORIZON); ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_TRIM_HORIZON);
Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap =
new HashMap<>(); new HashMap<>();
@ -877,10 +881,10 @@ public class ShardSyncerTest {
} }
assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size()); assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size());
for (KinesisClientLease lease : newLeases) { for (Lease lease : newLeases) {
assertTrue("Unexpected lease: " + lease, assertTrue("Unexpected lease: " + lease,
expectedShardIdCheckpointMap.containsKey(lease.getLeaseKey())); expectedShardIdCheckpointMap.containsKey(lease.leaseKey()));
assertEquals(expectedShardIdCheckpointMap.get(lease.getLeaseKey()), lease.getCheckpoint()); assertEquals(expectedShardIdCheckpointMap.get(lease.leaseKey()), lease.checkpoint());
} }
} }
@ -897,14 +901,14 @@ public class ShardSyncerTest {
@Test @Test
public final void testDetermineNewLeasesToCreateSplitMergeAtTimestamp1() { public final void testDetermineNewLeasesToCreateSplitMergeAtTimestamp1() {
List<Shard> shards = constructShardListForGraphA(); List<Shard> shards = constructShardListForGraphA();
List<KinesisClientLease> currentLeases = new ArrayList<>(); List<Lease> currentLeases = new ArrayList<>();
currentLeases.add(newLease("shardId-3")); currentLeases.add(newLease("shardId-3"));
currentLeases.add(newLease("shardId-4")); currentLeases.add(newLease("shardId-4"));
currentLeases.add(newLease("shardId-5")); currentLeases.add(newLease("shardId-5"));
List<KinesisClientLease> newLeases = List<Lease> newLeases =
ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_AT_TIMESTAMP); ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_AT_TIMESTAMP);
Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = new HashMap<>(); Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = new HashMap<>();
expectedShardIdCheckpointMap.put("shardId-8", ExtendedSequenceNumber.AT_TIMESTAMP); expectedShardIdCheckpointMap.put("shardId-8", ExtendedSequenceNumber.AT_TIMESTAMP);
@ -917,10 +921,10 @@ public class ShardSyncerTest {
expectedShardIdCheckpointMap.put("shardId-1", ExtendedSequenceNumber.AT_TIMESTAMP); expectedShardIdCheckpointMap.put("shardId-1", ExtendedSequenceNumber.AT_TIMESTAMP);
assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size()); assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size());
for (KinesisClientLease lease : newLeases) { for (Lease lease : newLeases) {
assertTrue("Unexpected lease: " + lease, assertTrue("Unexpected lease: " + lease,
expectedShardIdCheckpointMap.containsKey(lease.getLeaseKey())); expectedShardIdCheckpointMap.containsKey(lease.leaseKey()));
assertEquals(expectedShardIdCheckpointMap.get(lease.getLeaseKey()), lease.getCheckpoint()); assertEquals(expectedShardIdCheckpointMap.get(lease.leaseKey()), lease.checkpoint());
} }
} }
@ -937,13 +941,13 @@ public class ShardSyncerTest {
@Test @Test
public final void testDetermineNewLeasesToCreateSplitMergeAtTimestamp2() { public final void testDetermineNewLeasesToCreateSplitMergeAtTimestamp2() {
List<Shard> shards = constructShardListForGraphA(); List<Shard> shards = constructShardListForGraphA();
List<KinesisClientLease> currentLeases = new ArrayList<>(); List<Lease> currentLeases = new ArrayList<>();
currentLeases.add(newLease("shardId-4")); currentLeases.add(newLease("shardId-4"));
currentLeases.add(newLease("shardId-5")); currentLeases.add(newLease("shardId-5"));
currentLeases.add(newLease("shardId-7")); currentLeases.add(newLease("shardId-7"));
List<KinesisClientLease> newLeases = List<Lease> newLeases =
ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_AT_TIMESTAMP); ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_AT_TIMESTAMP);
Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = new HashMap<>(); Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = new HashMap<>();
expectedShardIdCheckpointMap.put("shardId-8", ExtendedSequenceNumber.AT_TIMESTAMP); expectedShardIdCheckpointMap.put("shardId-8", ExtendedSequenceNumber.AT_TIMESTAMP);
@ -954,10 +958,10 @@ public class ShardSyncerTest {
expectedShardIdCheckpointMap.put("shardId-1", ExtendedSequenceNumber.AT_TIMESTAMP); expectedShardIdCheckpointMap.put("shardId-1", ExtendedSequenceNumber.AT_TIMESTAMP);
assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size()); assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size());
for (KinesisClientLease lease : newLeases) { for (Lease lease : newLeases) {
assertTrue("Unexpected lease: " + lease, assertTrue("Unexpected lease: " + lease,
expectedShardIdCheckpointMap.containsKey(lease.getLeaseKey())); expectedShardIdCheckpointMap.containsKey(lease.leaseKey()));
assertEquals(expectedShardIdCheckpointMap.get(lease.getLeaseKey()), lease.getCheckpoint()); assertEquals(expectedShardIdCheckpointMap.get(lease.leaseKey()), lease.checkpoint());
} }
} }
@ -969,8 +973,8 @@ public class ShardSyncerTest {
@Test @Test
public final void testDetermineNewLeasesToCreateGraphBNoInitialLeasesAtTimestamp() { public final void testDetermineNewLeasesToCreateGraphBNoInitialLeasesAtTimestamp() {
List<Shard> shards = constructShardListForGraphB(); List<Shard> shards = constructShardListForGraphB();
List<KinesisClientLease> currentLeases = new ArrayList<>(); List<Lease> currentLeases = new ArrayList<>();
List<KinesisClientLease> newLeases = List<Lease> newLeases =
ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_AT_TIMESTAMP); ShardSyncer.determineNewLeasesToCreate(shards, currentLeases, INITIAL_POSITION_AT_TIMESTAMP);
Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap = Map<String, ExtendedSequenceNumber> expectedShardIdCheckpointMap =
new HashMap<>(); new HashMap<>();
@ -980,10 +984,10 @@ public class ShardSyncerTest {
} }
assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size()); assertEquals(expectedShardIdCheckpointMap.size(), newLeases.size());
for (KinesisClientLease lease : newLeases) { for (Lease lease : newLeases) {
assertTrue("Unexpected lease: " + lease, assertTrue("Unexpected lease: " + lease,
expectedShardIdCheckpointMap.containsKey(lease.getLeaseKey())); expectedShardIdCheckpointMap.containsKey(lease.leaseKey()));
assertEquals(expectedShardIdCheckpointMap.get(lease.getLeaseKey()), lease.getCheckpoint()); assertEquals(expectedShardIdCheckpointMap.get(lease.leaseKey()), lease.checkpoint());
} }
} }
@ -1112,7 +1116,7 @@ public class ShardSyncerTest {
kinesisShards.put(shardId, ShardObjectHelper.newShard(shardId, null, null, null)); kinesisShards.put(shardId, ShardObjectHelper.newShard(shardId, null, null, null));
Set<String> shardIdsOfCurrentLeases = new HashSet<String>(); Set<String> shardIdsOfCurrentLeases = new HashSet<String>();
shardIdsOfCurrentLeases.add(shardId); shardIdsOfCurrentLeases.add(shardId);
Map<String, KinesisClientLease> newLeaseMap = new HashMap<String, KinesisClientLease>(); Map<String, Lease> newLeaseMap = new HashMap<String, Lease>();
Map<String, Boolean> memoizationContext = new HashMap<>(); Map<String, Boolean> memoizationContext = new HashMap<>();
assertTrue(ShardSyncer.checkIfDescendantAndAddNewLeasesForAncestors(shardId, INITIAL_POSITION_LATEST, assertTrue(ShardSyncer.checkIfDescendantAndAddNewLeasesForAncestors(shardId, INITIAL_POSITION_LATEST,
shardIdsOfCurrentLeases, shardIdsOfCurrentLeases,
@ -1128,7 +1132,7 @@ public class ShardSyncerTest {
@Test @Test
public final void testCheckIfDescendantAndAddNewLeasesForAncestors2P2ANotDescendant() { public final void testCheckIfDescendantAndAddNewLeasesForAncestors2P2ANotDescendant() {
Set<String> shardIdsOfCurrentLeases = new HashSet<String>(); Set<String> shardIdsOfCurrentLeases = new HashSet<String>();
Map<String, KinesisClientLease> newLeaseMap = new HashMap<String, KinesisClientLease>(); Map<String, Lease> newLeaseMap = new HashMap<String, Lease>();
Map<String, Shard> kinesisShards = new HashMap<String, Shard>(); Map<String, Shard> kinesisShards = new HashMap<String, Shard>();
String parentShardId = "shardId-parent"; String parentShardId = "shardId-parent";
@ -1155,7 +1159,7 @@ public class ShardSyncerTest {
@Test @Test
public final void testCheckIfDescendantAndAddNewLeasesForAncestors2P2A1PDescendant() { public final void testCheckIfDescendantAndAddNewLeasesForAncestors2P2A1PDescendant() {
Set<String> shardIdsOfCurrentLeases = new HashSet<String>(); Set<String> shardIdsOfCurrentLeases = new HashSet<String>();
Map<String, KinesisClientLease> newLeaseMap = new HashMap<String, KinesisClientLease>(); Map<String, Lease> newLeaseMap = new HashMap<String, Lease>();
Map<String, Shard> kinesisShards = new HashMap<String, Shard>(); Map<String, Shard> kinesisShards = new HashMap<String, Shard>();
String parentShardId = "shardId-parent"; String parentShardId = "shardId-parent";
@ -1177,8 +1181,8 @@ public class ShardSyncerTest {
memoizationContext)); memoizationContext));
assertEquals(1, newLeaseMap.size()); assertEquals(1, newLeaseMap.size());
assertTrue(newLeaseMap.containsKey(adjacentParentShardId)); assertTrue(newLeaseMap.containsKey(adjacentParentShardId));
KinesisClientLease adjacentParentLease = newLeaseMap.get(adjacentParentShardId); Lease adjacentParentLease = newLeaseMap.get(adjacentParentShardId);
assertEquals(ExtendedSequenceNumber.LATEST, adjacentParentLease.getCheckpoint()); assertEquals(ExtendedSequenceNumber.LATEST, adjacentParentLease.checkpoint());
} }
/** /**
@ -1288,10 +1292,10 @@ public class ShardSyncerTest {
shard.setParentShardId(parentShardId); shard.setParentShardId(parentShardId);
shard.setAdjacentParentShardId(adjacentParentShardId); shard.setAdjacentParentShardId(adjacentParentShardId);
KinesisClientLease lease = ShardSyncer.newKCLLease(shard); Lease lease = ShardSyncer.newKCLLease(shard);
assertEquals(shardId, lease.getLeaseKey()); assertEquals(shardId, lease.leaseKey());
assertNull(lease.getCheckpoint()); assertNull(lease.checkpoint());
Set<String> parentIds = lease.getParentShardIds(); Set<String> parentIds = lease.parentShardIds();
assertEquals(2, parentIds.size()); assertEquals(2, parentIds.size());
assertTrue(parentIds.contains(parentShardId)); assertTrue(parentIds.contains(parentShardId));
assertTrue(parentIds.contains(adjacentParentShardId)); assertTrue(parentIds.contains(adjacentParentShardId));
@ -1364,11 +1368,11 @@ public class ShardSyncerTest {
String parentShardId = "shardId-0000"; String parentShardId = "shardId-0000";
String adjacentParentShardId = "shardId-0001"; String adjacentParentShardId = "shardId-0001";
String shardId = "shardId-0002"; String shardId = "shardId-0002";
KinesisClientLease lease = newLease(shardId); Lease lease = newLease(shardId);
List<String> parentShardIds = new ArrayList<>(); List<String> parentShardIds = new ArrayList<>();
parentShardIds.add(parentShardId); parentShardIds.add(parentShardId);
parentShardIds.add(adjacentParentShardId); parentShardIds.add(adjacentParentShardId);
lease.setParentShardIds(parentShardIds); lease.parentShardIds(parentShardIds);
Set<String> currentKinesisShardIds = new HashSet<>(); Set<String> currentKinesisShardIds = new HashSet<>();
currentKinesisShardIds.add(shardId); currentKinesisShardIds.add(shardId);
@ -1401,11 +1405,11 @@ public class ShardSyncerTest {
String parentShardId = "shardId-0000"; String parentShardId = "shardId-0000";
String adjacentParentShardId = "shardId-0001"; String adjacentParentShardId = "shardId-0001";
String shardId = "shardId-0002"; String shardId = "shardId-0002";
KinesisClientLease lease = newLease(shardId); Lease lease = newLease(shardId);
List<String> parentShardIds = new ArrayList<>(); List<String> parentShardIds = new ArrayList<>();
parentShardIds.add(parentShardId); parentShardIds.add(parentShardId);
parentShardIds.add(adjacentParentShardId); parentShardIds.add(adjacentParentShardId);
lease.setParentShardIds(parentShardIds); lease.parentShardIds(parentShardIds);
Set<String> currentKinesisShardIds = new HashSet<>(); Set<String> currentKinesisShardIds = new HashSet<>();
currentKinesisShardIds.add(parentShardId); currentKinesisShardIds.add(parentShardId);
@ -1422,11 +1426,11 @@ public class ShardSyncerTest {
String parentShardId = "shardId-0000"; String parentShardId = "shardId-0000";
String adjacentParentShardId = "shardId-0001"; String adjacentParentShardId = "shardId-0001";
String shardId = "shardId-0002"; String shardId = "shardId-0002";
KinesisClientLease lease = newLease(shardId); Lease lease = newLease(shardId);
List<String> parentShardIds = new ArrayList<>(); List<String> parentShardIds = new ArrayList<>();
parentShardIds.add(parentShardId); parentShardIds.add(parentShardId);
parentShardIds.add(adjacentParentShardId); parentShardIds.add(adjacentParentShardId);
lease.setParentShardIds(parentShardIds); lease.parentShardIds(parentShardIds);
Set<String> currentKinesisShardIds = new HashSet<>(); Set<String> currentKinesisShardIds = new HashSet<>();
currentKinesisShardIds.add(adjacentParentShardId); currentKinesisShardIds.add(adjacentParentShardId);
@ -1444,68 +1448,68 @@ public class ShardSyncerTest {
public final void testCleanupLeaseForClosedShard() public final void testCleanupLeaseForClosedShard()
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
String closedShardId = "shardId-2"; String closedShardId = "shardId-2";
KinesisClientLease leaseForClosedShard = newLease(closedShardId); Lease leaseForClosedShard = newLease(closedShardId);
leaseForClosedShard.setCheckpoint(new ExtendedSequenceNumber("1234")); leaseForClosedShard.checkpoint(new ExtendedSequenceNumber("1234"));
dynamoDBLeaseManager.createLeaseIfNotExists(leaseForClosedShard); dynamoDBLeaseRefresher.createLeaseIfNotExists(leaseForClosedShard);
Set<String> childShardIds = new HashSet<>(); Set<String> childShardIds = new HashSet<>();
List<KinesisClientLease> trackedLeases = new ArrayList<>(); List<Lease> trackedLeases = new ArrayList<>();
Set<String> parentShardIds = new HashSet<>(); Set<String> parentShardIds = new HashSet<>();
parentShardIds.add(closedShardId); parentShardIds.add(closedShardId);
String childShardId1 = "shardId-5"; String childShardId1 = "shardId-5";
KinesisClientLease childLease1 = newLease(childShardId1); Lease childLease1 = newLease(childShardId1);
childLease1.setParentShardIds(parentShardIds); childLease1.parentShardIds(parentShardIds);
childLease1.setCheckpoint(ExtendedSequenceNumber.TRIM_HORIZON); childLease1.checkpoint(ExtendedSequenceNumber.TRIM_HORIZON);
String childShardId2 = "shardId-7"; String childShardId2 = "shardId-7";
KinesisClientLease childLease2 = newLease(childShardId2); Lease childLease2 = newLease(childShardId2);
childLease2.setParentShardIds(parentShardIds); childLease2.parentShardIds(parentShardIds);
childLease2.setCheckpoint(ExtendedSequenceNumber.TRIM_HORIZON); childLease2.checkpoint(ExtendedSequenceNumber.TRIM_HORIZON);
Map<String, KinesisClientLease> trackedLeaseMap = ShardSyncer.constructShardIdToKCLLeaseMap(trackedLeases); Map<String, Lease> trackedLeaseMap = ShardSyncer.constructShardIdToKCLLeaseMap(trackedLeases);
// empty list of leases // empty list of leases
ShardSyncer.cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, dynamoDBLeaseManager); ShardSyncer.cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, dynamoDBLeaseRefresher);
assertNotNull(dynamoDBLeaseManager.getLease(closedShardId)); assertNotNull(dynamoDBLeaseRefresher.getLease(closedShardId));
// closed shard has not been fully processed yet (checkpoint != SHARD_END) // closed shard has not been fully processed yet (checkpoint != SHARD_END)
trackedLeases.add(leaseForClosedShard); trackedLeases.add(leaseForClosedShard);
trackedLeaseMap = ShardSyncer.constructShardIdToKCLLeaseMap(trackedLeases); trackedLeaseMap = ShardSyncer.constructShardIdToKCLLeaseMap(trackedLeases);
ShardSyncer.cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, dynamoDBLeaseManager); ShardSyncer.cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, dynamoDBLeaseRefresher);
assertNotNull(dynamoDBLeaseManager.getLease(closedShardId)); assertNotNull(dynamoDBLeaseRefresher.getLease(closedShardId));
// closed shard has been fully processed yet (checkpoint == SHARD_END) // closed shard has been fully processed yet (checkpoint == SHARD_END)
leaseForClosedShard.setCheckpoint(ExtendedSequenceNumber.SHARD_END); leaseForClosedShard.checkpoint(ExtendedSequenceNumber.SHARD_END);
dynamoDBLeaseManager.updateLease(leaseForClosedShard); dynamoDBLeaseRefresher.updateLease(leaseForClosedShard);
ShardSyncer.cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, dynamoDBLeaseManager); ShardSyncer.cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, dynamoDBLeaseRefresher);
assertNull(dynamoDBLeaseManager.getLease(closedShardId)); assertNull(dynamoDBLeaseRefresher.getLease(closedShardId));
// lease for only one child exists // lease for only one child exists
childShardIds.add(childShardId1); childShardIds.add(childShardId1);
childShardIds.add(childShardId2); childShardIds.add(childShardId2);
dynamoDBLeaseManager.createLeaseIfNotExists(leaseForClosedShard); dynamoDBLeaseRefresher.createLeaseIfNotExists(leaseForClosedShard);
dynamoDBLeaseManager.createLeaseIfNotExists(childLease1); dynamoDBLeaseRefresher.createLeaseIfNotExists(childLease1);
trackedLeases.add(childLease1); trackedLeases.add(childLease1);
trackedLeaseMap = ShardSyncer.constructShardIdToKCLLeaseMap(trackedLeases); trackedLeaseMap = ShardSyncer.constructShardIdToKCLLeaseMap(trackedLeases);
ShardSyncer.cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, dynamoDBLeaseManager); ShardSyncer.cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, dynamoDBLeaseRefresher);
assertNotNull(dynamoDBLeaseManager.getLease(closedShardId)); assertNotNull(dynamoDBLeaseRefresher.getLease(closedShardId));
// leases for both children exists, but they are both at TRIM_HORIZON // leases for both children exists, but they are both at TRIM_HORIZON
dynamoDBLeaseManager.createLeaseIfNotExists(childLease2); dynamoDBLeaseRefresher.createLeaseIfNotExists(childLease2);
trackedLeases.add(childLease2); trackedLeases.add(childLease2);
trackedLeaseMap = ShardSyncer.constructShardIdToKCLLeaseMap(trackedLeases); trackedLeaseMap = ShardSyncer.constructShardIdToKCLLeaseMap(trackedLeases);
ShardSyncer.cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, dynamoDBLeaseManager); ShardSyncer.cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, dynamoDBLeaseRefresher);
assertNotNull(dynamoDBLeaseManager.getLease(closedShardId)); assertNotNull(dynamoDBLeaseRefresher.getLease(closedShardId));
// leases for both children exists, one is at TRIM_HORIZON // leases for both children exists, one is at TRIM_HORIZON
childLease1.setCheckpoint(new ExtendedSequenceNumber("34890")); childLease1.checkpoint(new ExtendedSequenceNumber("34890"));
dynamoDBLeaseManager.updateLease(childLease1); dynamoDBLeaseRefresher.updateLease(childLease1);
ShardSyncer.cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, dynamoDBLeaseManager); ShardSyncer.cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, dynamoDBLeaseRefresher);
assertNotNull(dynamoDBLeaseManager.getLease(closedShardId)); assertNotNull(dynamoDBLeaseRefresher.getLease(closedShardId));
// leases for both children exists, NONE of them are at TRIM_HORIZON // leases for both children exists, NONE of them are at TRIM_HORIZON
childLease2.setCheckpoint(new ExtendedSequenceNumber("43789")); childLease2.checkpoint(new ExtendedSequenceNumber("43789"));
dynamoDBLeaseManager.updateLease(childLease2); dynamoDBLeaseRefresher.updateLease(childLease2);
ShardSyncer.cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, dynamoDBLeaseManager); ShardSyncer.cleanupLeaseForClosedShard(closedShardId, childShardIds, trackedLeaseMap, dynamoDBLeaseRefresher);
assertNull(dynamoDBLeaseManager.getLease(closedShardId)); assertNull(dynamoDBLeaseRefresher.getLease(closedShardId));
} }
/** /**
@ -1675,9 +1679,9 @@ public class ShardSyncerTest {
* @param shardId * @param shardId
* @return * @return
*/ */
private KinesisClientLease newLease(String shardId) { private Lease newLease(String shardId) {
KinesisClientLease lease = new KinesisClientLease(); Lease lease = new Lease();
lease.setLeaseKey(shardId); lease.leaseKey(shardId);
return lease; return lease;
} }

View file

@ -1,176 +0,0 @@
/*
* Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Amazon Software License (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/asl/
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*/
package software.amazon.kinesis.leases;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import org.junit.Assert;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.LeasingException;
public class TestHarnessBuilder {
private long currentTimeNanos;
private Map<String, KinesisClientLease> leases = new HashMap<String, KinesisClientLease>();
private KinesisClientDynamoDBLeaseManager leaseManager;
private Map<String, KinesisClientLease> originalLeases = new HashMap<>();
private Callable<Long> timeProvider = new Callable<Long>() {
@Override
public Long call() throws Exception {
return currentTimeNanos;
}
};
public TestHarnessBuilder(KinesisClientDynamoDBLeaseManager leaseManager) {
this.leaseManager = leaseManager;
}
public TestHarnessBuilder withLease(String shardId) {
return withLease(shardId, "leaseOwner");
}
public TestHarnessBuilder withLease(String shardId, String owner) {
KinesisClientLease lease = createLease(shardId, owner);
KinesisClientLease originalLease = createLease(shardId, owner);
leases.put(shardId, lease);
originalLeases.put(shardId, originalLease);
return this;
}
private KinesisClientLease createLease(String shardId, String owner) {
KinesisClientLease lease = new KinesisClientLease();
lease.setCheckpoint(new ExtendedSequenceNumber("checkpoint"));
lease.setOwnerSwitchesSinceCheckpoint(0L);
lease.setLeaseCounter(0L);
lease.setLeaseOwner(owner);
lease.setParentShardIds(Collections.singleton("parentShardId"));
lease.setLeaseKey(shardId);
return lease;
}
public Map<String, KinesisClientLease> build() throws LeasingException {
for (KinesisClientLease lease : leases.values()) {
leaseManager.createLeaseIfNotExists(lease);
if (lease.getLeaseOwner() != null) {
lease.setLastCounterIncrementNanos(System.nanoTime());
}
}
currentTimeNanos = System.nanoTime();
return leases;
}
public void passTime(long millis) {
currentTimeNanos += millis * 1000000;
}
public Map<String, KinesisClientLease> takeMutateAssert(DynamoDBLeaseTaker<KinesisClientLease> taker, int numToTake)
throws LeasingException {
Map<String, KinesisClientLease> result = taker.takeLeases(timeProvider);
Assert.assertEquals(numToTake, result.size());
for (KinesisClientLease actual : result.values()) {
KinesisClientLease original = leases.get(actual.getLeaseKey());
Assert.assertNotNull(original);
mutateAssert(taker.getWorkerIdentifier(), original, actual);
}
return result;
}
public Map<String, KinesisClientLease> takeMutateAssert(DynamoDBLeaseTaker<KinesisClientLease> taker, String... takenShardIds)
throws LeasingException {
Map<String, KinesisClientLease> result = taker.takeLeases(timeProvider);
Assert.assertEquals(takenShardIds.length, result.size());
for (String shardId : takenShardIds) {
KinesisClientLease original = leases.get(shardId);
Assert.assertNotNull(original);
KinesisClientLease actual = result.get(shardId);
Assert.assertNotNull(actual);
mutateAssert(taker.getWorkerIdentifier(), original, actual);
}
return result;
}
private void mutateAssert(String newWorkerIdentifier, KinesisClientLease original, KinesisClientLease actual) {
original.setLeaseCounter(original.getLeaseCounter() + 1);
if (original.getLeaseOwner() != null && !newWorkerIdentifier.equals(original.getLeaseOwner())) {
original.setOwnerSwitchesSinceCheckpoint(original.getOwnerSwitchesSinceCheckpoint() + 1);
}
original.setLeaseOwner(newWorkerIdentifier);
Assert.assertEquals(original, actual); // Assert the contents of the lease
}
public void addLeasesToRenew(LeaseRenewer<KinesisClientLease> renewer, String... shardIds)
throws DependencyException, InvalidStateException {
List<KinesisClientLease> leasesToRenew = new ArrayList<KinesisClientLease>();
for (String shardId : shardIds) {
KinesisClientLease lease = leases.get(shardId);
Assert.assertNotNull(lease);
leasesToRenew.add(lease);
}
renewer.addLeasesToRenew(leasesToRenew);
}
public Map<String, KinesisClientLease> renewMutateAssert(LeaseRenewer<KinesisClientLease> renewer, String... renewedShardIds)
throws DependencyException, InvalidStateException {
renewer.renewLeases();
Map<String, KinesisClientLease> heldLeases = renewer.getCurrentlyHeldLeases();
Assert.assertEquals(renewedShardIds.length, heldLeases.size());
for (String shardId : renewedShardIds) {
KinesisClientLease original = originalLeases.get(shardId);
Assert.assertNotNull(original);
KinesisClientLease actual = heldLeases.get(shardId);
Assert.assertNotNull(actual);
original.setLeaseCounter(original.getLeaseCounter() + 1);
Assert.assertEquals(original, actual);
}
return heldLeases;
}
public void renewAllLeases() throws LeasingException {
for (KinesisClientLease lease : leases.values()) {
leaseManager.renewLease(lease);
}
}
}

View file

@ -12,7 +12,7 @@
* express or implied. See the License for the specific language governing * express or implied. See the License for the specific language governing
* permissions and limitations under the License. * permissions and limitations under the License.
*/ */
package software.amazon.kinesis.leases; package software.amazon.kinesis.leases.dynamodb;
import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertFalse;
@ -36,6 +36,9 @@ import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient; import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient;
import software.amazon.kinesis.checkpoint.DynamoDBCheckpointer; import software.amazon.kinesis.checkpoint.DynamoDBCheckpointer;
import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.LeaseCoordinator;
import software.amazon.kinesis.leases.LeaseRenewer;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.LeasingException; import software.amazon.kinesis.leases.exceptions.LeasingException;
@ -45,35 +48,35 @@ import software.amazon.kinesis.metrics.NullMetricsFactory;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber; import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
@RunWith(MockitoJUnitRunner.class) @RunWith(MockitoJUnitRunner.class)
public class KinesisClientLibLeaseCoordinatorIntegrationTest { public class DynamoDBLeaseCoordinatorIntegrationTest {
private static final String TABLE_NAME = KinesisClientLibLeaseCoordinatorIntegrationTest.class.getSimpleName(); private static final String TABLE_NAME = DynamoDBLeaseCoordinatorIntegrationTest.class.getSimpleName();
private static final String WORKER_ID = UUID.randomUUID().toString(); private static final String WORKER_ID = UUID.randomUUID().toString();
private static final long LEASE_DURATION_MILLIS = 5000L; private static final long LEASE_DURATION_MILLIS = 5000L;
private static final long EPSILON_MILLIS = 25L; private static final long EPSILON_MILLIS = 25L;
private static final int MAX_LEASES_FOR_WORKER = Integer.MAX_VALUE; private static final int MAX_LEASES_FOR_WORKER = Integer.MAX_VALUE;
private static final int MAX_LEASES_TO_STEAL_AT_ONE_TIME = 1; private static final int MAX_LEASES_TO_STEAL_AT_ONE_TIME = 1;
private static final int MAX_LEASE_RENEWER_THREAD_COUNT = 20; private static final int MAX_LEASE_RENEWER_THREAD_COUNT = 20;
private static KinesisClientDynamoDBLeaseManager leaseManager; private static DynamoDBLeaseRefresher leaseRefresher;
private static DynamoDBCheckpointer dynamoDBCheckpointer; private static DynamoDBCheckpointer dynamoDBCheckpointer;
private KinesisClientLibLeaseCoordinator coordinator; private LeaseCoordinator coordinator;
private final String leaseKey = "shd-1"; private final String leaseKey = "shd-1";
private final IMetricsFactory metricsFactory = new NullMetricsFactory(); private final IMetricsFactory metricsFactory = new NullMetricsFactory();
@Before @Before
public void setUp() throws ProvisionedThroughputException, DependencyException, InvalidStateException { public void setUp() throws ProvisionedThroughputException, DependencyException, InvalidStateException {
final boolean useConsistentReads = true; final boolean useConsistentReads = true;
if (leaseManager == null) { if (leaseRefresher == null) {
AmazonDynamoDBClient ddb = new AmazonDynamoDBClient(new DefaultAWSCredentialsProviderChain()); AmazonDynamoDBClient ddb = new AmazonDynamoDBClient(new DefaultAWSCredentialsProviderChain());
leaseManager = leaseRefresher =
new KinesisClientDynamoDBLeaseManager(TABLE_NAME, ddb, useConsistentReads); new DynamoDBLeaseRefresher(TABLE_NAME, ddb, new DynamoDBLeaseSerializer(), useConsistentReads);
} }
leaseManager.createLeaseTableIfNotExists(10L, 10L); leaseRefresher.createLeaseTableIfNotExists(10L, 10L);
leaseManager.deleteAll(); leaseRefresher.deleteAll();
coordinator = new KinesisClientLibLeaseCoordinator(leaseManager, WORKER_ID, LEASE_DURATION_MILLIS, coordinator = new DynamoDBLeaseCoordinator(leaseRefresher, WORKER_ID, LEASE_DURATION_MILLIS,
EPSILON_MILLIS, MAX_LEASES_FOR_WORKER, MAX_LEASES_TO_STEAL_AT_ONE_TIME, MAX_LEASE_RENEWER_THREAD_COUNT, EPSILON_MILLIS, MAX_LEASES_FOR_WORKER, MAX_LEASES_TO_STEAL_AT_ONE_TIME, MAX_LEASE_RENEWER_THREAD_COUNT,
metricsFactory); metricsFactory);
dynamoDBCheckpointer = new DynamoDBCheckpointer(coordinator, leaseManager, metricsFactory); dynamoDBCheckpointer = new DynamoDBCheckpointer(coordinator, leaseRefresher, metricsFactory);
coordinator.start(); coordinator.start();
} }
@ -86,14 +89,14 @@ public class KinesisClientLibLeaseCoordinatorIntegrationTest {
TestHarnessBuilder builder = new TestHarnessBuilder(); TestHarnessBuilder builder = new TestHarnessBuilder();
builder.withLease(leaseKey, null).build(); builder.withLease(leaseKey, null).build();
// Run the taker and renewer in-between getting the Lease object and calling setCheckpoint // Run the taker and renewer in-between getting the Lease object and calling checkpoint
coordinator.runLeaseTaker(); coordinator.runLeaseTaker();
coordinator.runLeaseRenewer(); coordinator.runLeaseRenewer();
KinesisClientLease lease = coordinator.getCurrentlyHeldLease(leaseKey); Lease lease = coordinator.getCurrentlyHeldLease(leaseKey);
if (lease == null) { if (lease == null) {
List<KinesisClientLease> leases = leaseManager.listLeases(); List<Lease> leases = leaseRefresher.listLeases();
for (KinesisClientLease kinesisClientLease : leases) { for (Lease kinesisClientLease : leases) {
System.out.println(kinesisClientLease); System.out.println(kinesisClientLease);
} }
} }
@ -101,13 +104,13 @@ public class KinesisClientLibLeaseCoordinatorIntegrationTest {
assertNotNull(lease); assertNotNull(lease);
ExtendedSequenceNumber newCheckpoint = new ExtendedSequenceNumber("newCheckpoint"); ExtendedSequenceNumber newCheckpoint = new ExtendedSequenceNumber("newCheckpoint");
// lease's leaseCounter is wrong at this point, but it shouldn't matter. // lease's leaseCounter is wrong at this point, but it shouldn't matter.
assertTrue(dynamoDBCheckpointer.setCheckpoint(lease.getLeaseKey(), newCheckpoint, lease.getConcurrencyToken())); assertTrue(dynamoDBCheckpointer.setCheckpoint(lease.leaseKey(), newCheckpoint, lease.concurrencyToken()));
Lease fromDynamo = leaseManager.getLease(lease.getLeaseKey()); Lease fromDynamo = leaseRefresher.getLease(lease.leaseKey());
lease.setLeaseCounter(lease.getLeaseCounter() + 1); lease.leaseCounter(lease.leaseCounter() + 1);
lease.setCheckpoint(newCheckpoint); lease.checkpoint(newCheckpoint);
lease.setLeaseOwner(coordinator.getWorkerIdentifier()); lease.leaseOwner(coordinator.workerIdentifier());
assertEquals(lease, fromDynamo); assertEquals(lease, fromDynamo);
} }
@ -121,19 +124,19 @@ public class KinesisClientLibLeaseCoordinatorIntegrationTest {
coordinator.runLeaseTaker(); coordinator.runLeaseTaker();
coordinator.runLeaseRenewer(); coordinator.runLeaseRenewer();
KinesisClientLease lease = coordinator.getCurrentlyHeldLease(leaseKey); Lease lease = coordinator.getCurrentlyHeldLease(leaseKey);
assertNotNull(lease); assertNotNull(lease);
leaseManager.renewLease(coordinator.getCurrentlyHeldLease(leaseKey)); leaseRefresher.renewLease(coordinator.getCurrentlyHeldLease(leaseKey));
ExtendedSequenceNumber newCheckpoint = new ExtendedSequenceNumber("newCheckpoint"); ExtendedSequenceNumber newCheckpoint = new ExtendedSequenceNumber("newCheckpoint");
assertFalse(dynamoDBCheckpointer.setCheckpoint(lease.getLeaseKey(), newCheckpoint, lease.getConcurrencyToken())); assertFalse(dynamoDBCheckpointer.setCheckpoint(lease.leaseKey(), newCheckpoint, lease.concurrencyToken()));
Lease fromDynamo = leaseManager.getLease(lease.getLeaseKey()); Lease fromDynamo = leaseRefresher.getLease(lease.leaseKey());
lease.setLeaseCounter(lease.getLeaseCounter() + 1); lease.leaseCounter(lease.leaseCounter() + 1);
// Counter and owner changed, but checkpoint did not. // Counter and owner changed, but checkpoint did not.
lease.setLeaseOwner(coordinator.getWorkerIdentifier()); lease.leaseOwner(coordinator.workerIdentifier());
assertEquals(lease, fromDynamo); assertEquals(lease, fromDynamo);
} }
@ -147,17 +150,17 @@ public class KinesisClientLibLeaseCoordinatorIntegrationTest {
coordinator.runLeaseTaker(); coordinator.runLeaseTaker();
coordinator.runLeaseRenewer(); coordinator.runLeaseRenewer();
KinesisClientLease lease = coordinator.getCurrentlyHeldLease(leaseKey); Lease lease = coordinator.getCurrentlyHeldLease(leaseKey);
assertNotNull(lease); assertNotNull(lease);
ExtendedSequenceNumber newCheckpoint = new ExtendedSequenceNumber("newCheckpoint"); ExtendedSequenceNumber newCheckpoint = new ExtendedSequenceNumber("newCheckpoint");
assertFalse(dynamoDBCheckpointer.setCheckpoint(lease.getLeaseKey(), newCheckpoint, UUID.randomUUID())); assertFalse(dynamoDBCheckpointer.setCheckpoint(lease.leaseKey(), newCheckpoint, UUID.randomUUID()));
Lease fromDynamo = leaseManager.getLease(lease.getLeaseKey()); Lease fromDynamo = leaseRefresher.getLease(lease.leaseKey());
// Owner should be the only thing that changed. // Owner should be the only thing that changed.
lease.setLeaseOwner(coordinator.getWorkerIdentifier()); lease.leaseOwner(coordinator.workerIdentifier());
assertEquals(lease, fromDynamo); assertEquals(lease, fromDynamo);
} }
@ -165,7 +168,7 @@ public class KinesisClientLibLeaseCoordinatorIntegrationTest {
private long currentTimeNanos; private long currentTimeNanos;
private Map<String, KinesisClientLease> leases = new HashMap<String, KinesisClientLease>(); private Map<String, Lease> leases = new HashMap<String, Lease>();
private Callable<Long> timeProvider = new Callable<Long>() { private Callable<Long> timeProvider = new Callable<Long>() {
@ -181,23 +184,23 @@ public class KinesisClientLibLeaseCoordinatorIntegrationTest {
} }
public TestHarnessBuilder withLease(String shardId, String owner) { public TestHarnessBuilder withLease(String shardId, String owner) {
KinesisClientLease lease = new KinesisClientLease(); Lease lease = new Lease();
lease.setCheckpoint(new ExtendedSequenceNumber("checkpoint")); lease.checkpoint(new ExtendedSequenceNumber("checkpoint"));
lease.setOwnerSwitchesSinceCheckpoint(0L); lease.ownerSwitchesSinceCheckpoint(0L);
lease.setLeaseCounter(0L); lease.leaseCounter(0L);
lease.setLeaseOwner(owner); lease.leaseOwner(owner);
lease.setParentShardIds(Collections.singleton("parentShardId")); lease.parentShardIds(Collections.singleton("parentShardId"));
lease.setLeaseKey(shardId); lease.leaseKey(shardId);
leases.put(shardId, lease); leases.put(shardId, lease);
return this; return this;
} }
public Map<String, KinesisClientLease> build() throws LeasingException { public Map<String, Lease> build() throws LeasingException {
for (KinesisClientLease lease : leases.values()) { for (Lease lease : leases.values()) {
leaseManager.createLeaseIfNotExists(lease); leaseRefresher.createLeaseIfNotExists(lease);
if (lease.getLeaseOwner() != null) { if (lease.leaseOwner() != null) {
lease.setLastCounterIncrementNanos(System.nanoTime()); lease.lastCounterIncrementNanos(System.nanoTime());
} }
} }
@ -210,22 +213,22 @@ public class KinesisClientLibLeaseCoordinatorIntegrationTest {
currentTimeNanos += millis * 1000000; currentTimeNanos += millis * 1000000;
} }
private void mutateAssert(String newWorkerIdentifier, KinesisClientLease original, KinesisClientLease actual) { private void mutateAssert(String newWorkerIdentifier, Lease original, Lease actual) {
original.setLeaseCounter(original.getLeaseCounter() + 1); original.leaseCounter(original.leaseCounter() + 1);
if (original.getLeaseOwner() != null && !newWorkerIdentifier.equals(original.getLeaseOwner())) { if (original.leaseOwner() != null && !newWorkerIdentifier.equals(original.leaseOwner())) {
original.setOwnerSwitchesSinceCheckpoint(original.getOwnerSwitchesSinceCheckpoint() + 1); original.ownerSwitchesSinceCheckpoint(original.ownerSwitchesSinceCheckpoint() + 1);
} }
original.setLeaseOwner(newWorkerIdentifier); original.leaseOwner(newWorkerIdentifier);
assertEquals(original, actual); // Assert the contents of the lease assertEquals(original, actual); // Assert the contents of the lease
} }
public void addLeasesToRenew(LeaseRenewer<KinesisClientLease> renewer, String... shardIds) public void addLeasesToRenew(LeaseRenewer renewer, String... shardIds)
throws DependencyException, InvalidStateException { throws DependencyException, InvalidStateException {
List<KinesisClientLease> leasesToRenew = new ArrayList<KinesisClientLease>(); List<Lease> leasesToRenew = new ArrayList<>();
for (String shardId : shardIds) { for (String shardId : shardIds) {
KinesisClientLease lease = leases.get(shardId); Lease lease = leases.get(shardId);
assertNotNull(lease); assertNotNull(lease);
leasesToRenew.add(lease); leasesToRenew.add(lease);
} }
@ -233,21 +236,21 @@ public class KinesisClientLibLeaseCoordinatorIntegrationTest {
renewer.addLeasesToRenew(leasesToRenew); renewer.addLeasesToRenew(leasesToRenew);
} }
public Map<String, KinesisClientLease> renewMutateAssert(LeaseRenewer<KinesisClientLease> renewer, public Map<String, Lease> renewMutateAssert(LeaseRenewer renewer,
String... renewedShardIds) throws DependencyException, InvalidStateException { String... renewedShardIds) throws DependencyException, InvalidStateException {
renewer.renewLeases(); renewer.renewLeases();
Map<String, KinesisClientLease> heldLeases = renewer.getCurrentlyHeldLeases(); Map<String, Lease> heldLeases = renewer.getCurrentlyHeldLeases();
assertEquals(renewedShardIds.length, heldLeases.size()); assertEquals(renewedShardIds.length, heldLeases.size());
for (String shardId : renewedShardIds) { for (String shardId : renewedShardIds) {
KinesisClientLease original = leases.get(shardId); Lease original = leases.get(shardId);
assertNotNull(original); assertNotNull(original);
KinesisClientLease actual = heldLeases.get(shardId); Lease actual = heldLeases.get(shardId);
assertNotNull(actual); assertNotNull(actual);
original.setLeaseCounter(original.getLeaseCounter() + 1); original.leaseCounter(original.leaseCounter() + 1);
assertEquals(original, actual); assertEquals(original, actual);
} }
@ -255,8 +258,8 @@ public class KinesisClientLibLeaseCoordinatorIntegrationTest {
} }
public void renewAllLeases() throws LeasingException { public void renewAllLeases() throws LeasingException {
for (KinesisClientLease lease : leases.values()) { for (Lease lease : leases.values()) {
leaseManager.renewLease(lease); leaseRefresher.renewLease(lease);
} }
} }
} }

View file

@ -12,7 +12,7 @@
* express or implied. See the License for the specific language governing * express or implied. See the License for the specific language governing
* permissions and limitations under the License. * permissions and limitations under the License.
*/ */
package software.amazon.kinesis.leases; package software.amazon.kinesis.leases.dynamodb;
import static org.mockito.Matchers.eq; import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verify;
@ -24,29 +24,31 @@ import org.junit.Before;
import org.junit.Test; import org.junit.Test;
import org.junit.runner.RunWith; import org.junit.runner.RunWith;
import org.mockito.Mock; import org.mockito.Mock;
import org.mockito.runners.MockitoJUnitRunner;
import com.amazonaws.services.kinesis.clientlibrary.exceptions.KinesisClientLibException; import com.amazonaws.services.kinesis.clientlibrary.exceptions.KinesisClientLibException;
import com.amazonaws.services.kinesis.clientlibrary.exceptions.ShutdownException; import com.amazonaws.services.kinesis.clientlibrary.exceptions.ShutdownException;
import org.mockito.runners.MockitoJUnitRunner;
import software.amazon.kinesis.checkpoint.DynamoDBCheckpointer; import software.amazon.kinesis.checkpoint.DynamoDBCheckpointer;
import software.amazon.kinesis.metrics.IMetricsFactory; import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber; import software.amazon.kinesis.leases.LeaseCoordinator;
import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException; import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
import software.amazon.kinesis.metrics.IMetricsFactory;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
@RunWith(MockitoJUnitRunner.class) @RunWith(MockitoJUnitRunner.class)
public class KinesisClientLibLeaseCoordinatorTest { public class DynamoDBLeaseCoordinatorTest {
private static final String SHARD_ID = "shardId-test"; private static final String SHARD_ID = "shardId-test";
private static final String WORK_ID = "workId-test";
private static final long TEST_LONG = 1000L;
private static final ExtendedSequenceNumber TEST_CHKPT = new ExtendedSequenceNumber("string-test"); private static final ExtendedSequenceNumber TEST_CHKPT = new ExtendedSequenceNumber("string-test");
private static final UUID TEST_UUID = UUID.randomUUID(); private static final UUID TEST_UUID = UUID.randomUUID();
@Mock @Mock
private LeaseManager<KinesisClientLease> leaseManager; private LeaseRefresher leaseRefresher;
@Mock @Mock
private LeaseCoordinator<KinesisClientLease> leaseCoordinator; private LeaseCoordinator leaseCoordinator;
@Mock @Mock
private IMetricsFactory metricsFactory; private IMetricsFactory metricsFactory;
@ -55,18 +57,18 @@ public class KinesisClientLibLeaseCoordinatorTest {
@SuppressWarnings("unchecked") @SuppressWarnings("unchecked")
@Before @Before
public void setUpLeaseCoordinator() throws ProvisionedThroughputException, DependencyException { public void setUpLeaseCoordinator() throws ProvisionedThroughputException, DependencyException {
dynamoDBCheckpointer = new DynamoDBCheckpointer(leaseCoordinator, leaseManager, metricsFactory); dynamoDBCheckpointer = new DynamoDBCheckpointer(leaseCoordinator, leaseRefresher, metricsFactory);
} }
@Test(expected = ShutdownException.class) @Test(expected = ShutdownException.class)
public void testSetCheckpointWithUnownedShardId() throws KinesisClientLibException, DependencyException, InvalidStateException, ProvisionedThroughputException { public void testSetCheckpointWithUnownedShardId() throws KinesisClientLibException, DependencyException, InvalidStateException, ProvisionedThroughputException {
final KinesisClientLease lease = new KinesisClientLease(); final Lease lease = new Lease();
when(leaseManager.getLease(eq(SHARD_ID))).thenReturn(lease); when(leaseRefresher.getLease(eq(SHARD_ID))).thenReturn(lease);
when(leaseCoordinator.updateLease(eq(lease), eq(TEST_UUID))).thenReturn(false); when(leaseCoordinator.updateLease(eq(lease), eq(TEST_UUID))).thenReturn(false);
dynamoDBCheckpointer.setCheckpoint(SHARD_ID, TEST_CHKPT, TEST_UUID.toString()); dynamoDBCheckpointer.setCheckpoint(SHARD_ID, TEST_CHKPT, TEST_UUID.toString());
verify(leaseManager).getLease(eq(SHARD_ID)); verify(leaseRefresher).getLease(eq(SHARD_ID));
verify(leaseCoordinator).updateLease(eq(lease), eq(TEST_UUID)); verify(leaseCoordinator).updateLease(eq(lease), eq(TEST_UUID));
} }
@ -74,7 +76,7 @@ public class KinesisClientLibLeaseCoordinatorTest {
// public void testWaitLeaseTableTimeout() // public void testWaitLeaseTableTimeout()
// throws DependencyException, ProvisionedThroughputException, IllegalStateException { // throws DependencyException, ProvisionedThroughputException, IllegalStateException {
// Set mock lease manager to return false in waiting // Set mock lease manager to return false in waiting
// doReturn(false).when(leaseManager).waitUntilLeaseTableExists(anyLong(), anyLong()); // doReturn(false).when(leaseRefresher).waitUntilLeaseTableExists(anyLong(), anyLong());
// leaseCoordinator.initialize(); // leaseCoordinator.initialize();
// } // }
} }

View file

@ -12,27 +12,37 @@
* express or implied. See the License for the specific language governing * express or implied. See the License for the specific language governing
* permissions and limitations under the License. * permissions and limitations under the License.
*/ */
package software.amazon.kinesis.leases; package software.amazon.kinesis.leases.dynamodb;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import java.util.Collection; import java.util.Collection;
import java.util.List; import java.util.List;
import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicInteger;
import junit.framework.Assert;
import org.junit.Test; import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.runners.MockitoJUnitRunner;
import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.LeaseIntegrationTest;
import software.amazon.kinesis.leases.LeaseSerializer;
import software.amazon.kinesis.leases.exceptions.LeasingException; import software.amazon.kinesis.leases.exceptions.LeasingException;
public class DynamoDBLeaseManagerIntegrationTest extends LeaseIntegrationTest { public class DynamoDBLeaseRefresherIntegrationTest extends LeaseIntegrationTest {
/** /**
* Test listLeases when no records are present. * Test listLeases when no records are present.
*/ */
@Test @Test
public void testListNoRecords() throws LeasingException { public void testListNoRecords() throws LeasingException {
List<KinesisClientLease> leases = leaseManager.listLeases(); List<Lease> leases = leaseRefresher.listLeases();
Assert.assertTrue(leases.isEmpty()); assertTrue(leases.isEmpty());
} }
/** /**
@ -40,7 +50,7 @@ public class DynamoDBLeaseManagerIntegrationTest extends LeaseIntegrationTest {
*/ */
@Test @Test
public void testListWithRecords() throws LeasingException { public void testListWithRecords() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
int numRecordsToPut = 10; int numRecordsToPut = 10;
@ -48,16 +58,16 @@ public class DynamoDBLeaseManagerIntegrationTest extends LeaseIntegrationTest {
builder.withLease(Integer.toString(i)); builder.withLease(Integer.toString(i));
} }
Collection<KinesisClientLease> expected = builder.build().values(); Collection<Lease> expected = builder.build().values();
// The / 3 here ensures that we will test Dynamo's paging mechanics. // The / 3 here ensures that we will test Dynamo's paging mechanics.
List<KinesisClientLease> actual = leaseManager.list(numRecordsToPut / 3); List<Lease> actual = leaseRefresher.list(numRecordsToPut / 3);
for (KinesisClientLease lease : actual) { for (Lease lease : actual) {
Assert.assertNotNull(expected.remove(lease)); assertNotNull(expected.remove(lease));
} }
Assert.assertTrue(expected.isEmpty()); assertTrue(expected.isEmpty());
} }
/** /**
@ -65,53 +75,53 @@ public class DynamoDBLeaseManagerIntegrationTest extends LeaseIntegrationTest {
*/ */
@Test @Test
public void testGetLease() throws LeasingException { public void testGetLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
Lease expected = builder.withLease("1").build().get("1"); Lease expected = builder.withLease("1").build().get("1");
Lease actual = leaseManager.getLease(expected.getLeaseKey()); Lease actual = leaseRefresher.getLease(expected.leaseKey());
Assert.assertEquals(expected, actual); assertEquals(expected, actual);
} }
/** /**
* Tests leaseManager.get() when the looked-for record is absent. * Tests leaseRefresher.get() when the looked-for record is absent.
*/ */
@Test @Test
public void testGetNull() throws LeasingException { public void testGetNull() throws LeasingException {
Lease actual = leaseManager.getLease("bogusShardId"); Lease actual = leaseRefresher.getLease("bogusShardId");
Assert.assertNull(actual); assertNull(actual);
} }
/** /**
* Tests leaseManager.holdLease's success scenario. * Tests leaseRefresher.holdLease's success scenario.
*/ */
@Test @Test
public void testRenewLease() throws LeasingException { public void testRenewLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
KinesisClientLease lease = builder.withLease("1").build().get("1"); Lease lease = builder.withLease("1").build().get("1");
Long originalLeaseCounter = lease.getLeaseCounter(); Long originalLeaseCounter = lease.leaseCounter();
leaseManager.renewLease(lease); leaseRefresher.renewLease(lease);
Assert.assertTrue(originalLeaseCounter + 1 == lease.getLeaseCounter()); assertTrue(originalLeaseCounter + 1 == lease.leaseCounter());
Lease fromDynamo = leaseManager.getLease(lease.getLeaseKey()); Lease fromDynamo = leaseRefresher.getLease(lease.leaseKey());
Assert.assertEquals(lease, fromDynamo); assertEquals(lease, fromDynamo);
} }
/** /**
* Tests leaseManager.holdLease when the lease has changed out from under us. * Tests leaseRefresher.holdLease when the lease has changed out from under us.
*/ */
@Test @Test
public void testHoldUpdatedLease() throws LeasingException { public void testHoldUpdatedLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
KinesisClientLease lease = builder.withLease("1").build().get("1"); Lease lease = builder.withLease("1").build().get("1");
KinesisClientLease leaseCopy = leaseManager.getLease(lease.getLeaseKey()); Lease leaseCopy = leaseRefresher.getLease(lease.leaseKey());
// lose lease // lose lease
leaseManager.takeLease(lease, "bar"); leaseRefresher.takeLease(lease, "bar");
Assert.assertFalse(leaseManager.renewLease(leaseCopy)); assertFalse(leaseRefresher.renewLease(leaseCopy));
} }
/** /**
@ -131,19 +141,19 @@ public class DynamoDBLeaseManagerIntegrationTest extends LeaseIntegrationTest {
} }
private void testTakeLease(boolean owned) throws LeasingException { private void testTakeLease(boolean owned) throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
KinesisClientLease lease = builder.withLease("1", owned ? "originalOwner" : null).build().get("1"); Lease lease = builder.withLease("1", owned ? "originalOwner" : null).build().get("1");
Long originalLeaseCounter = lease.getLeaseCounter(); Long originalLeaseCounter = lease.leaseCounter();
String newOwner = "newOwner"; String newOwner = "newOwner";
leaseManager.takeLease(lease, newOwner); leaseRefresher.takeLease(lease, newOwner);
Assert.assertTrue(originalLeaseCounter + 1 == lease.getLeaseCounter()); assertTrue(originalLeaseCounter + 1 == lease.leaseCounter());
Assert.assertTrue((owned ? 1 : 0) == lease.getOwnerSwitchesSinceCheckpoint()); assertTrue((owned ? 1 : 0) == lease.ownerSwitchesSinceCheckpoint());
Assert.assertEquals(newOwner, lease.getLeaseOwner()); assertEquals(newOwner, lease.leaseOwner());
Lease fromDynamo = leaseManager.getLease(lease.getLeaseKey()); Lease fromDynamo = leaseRefresher.getLease(lease.leaseKey());
Assert.assertEquals(lease, fromDynamo); assertEquals(lease, fromDynamo);
} }
/** /**
@ -151,25 +161,25 @@ public class DynamoDBLeaseManagerIntegrationTest extends LeaseIntegrationTest {
*/ */
@Test @Test
public void testTakeUpdatedLease() throws LeasingException { public void testTakeUpdatedLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
KinesisClientLease lease = builder.withLease("1").build().get("1"); Lease lease = builder.withLease("1").build().get("1");
KinesisClientLease leaseCopy = leaseManager.getLease(lease.getLeaseKey()); Lease leaseCopy = leaseRefresher.getLease(lease.leaseKey());
String newOwner = "newOwner"; String newOwner = "newOwner";
leaseManager.takeLease(lease, newOwner); leaseRefresher.takeLease(lease, newOwner);
Assert.assertFalse(leaseManager.takeLease(leaseCopy, newOwner)); assertFalse(leaseRefresher.takeLease(leaseCopy, newOwner));
} }
/** /**
* Tests evictLease when the lease is currently unowned. * Tests evictLease when the lease is currently unowned.
*/ */
public void testEvictUnownedLease() throws LeasingException { public void testEvictUnownedLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
KinesisClientLease lease = builder.withLease("1", null).build().get("1"); Lease lease = builder.withLease("1", null).build().get("1");
Assert.assertFalse(leaseManager.evictLease(lease)); assertFalse(leaseRefresher.evictLease(lease));
} }
/** /**
@ -177,17 +187,17 @@ public class DynamoDBLeaseManagerIntegrationTest extends LeaseIntegrationTest {
*/ */
@Test @Test
public void testEvictOwnedLease() throws LeasingException { public void testEvictOwnedLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
KinesisClientLease lease = builder.withLease("1").build().get("1"); Lease lease = builder.withLease("1").build().get("1");
Long originalLeaseCounter = lease.getLeaseCounter(); Long originalLeaseCounter = lease.leaseCounter();
leaseManager.evictLease(lease); leaseRefresher.evictLease(lease);
Assert.assertNull(lease.getLeaseOwner()); assertNull(lease.leaseOwner());
Assert.assertTrue(originalLeaseCounter + 1 == lease.getLeaseCounter()); assertTrue(originalLeaseCounter + 1 == lease.leaseCounter());
Lease fromDynamo = leaseManager.getLease(lease.getLeaseKey()); Lease fromDynamo = leaseRefresher.getLease(lease.leaseKey());
Assert.assertEquals(lease, fromDynamo); assertEquals(lease, fromDynamo);
} }
/** /**
@ -197,12 +207,12 @@ public class DynamoDBLeaseManagerIntegrationTest extends LeaseIntegrationTest {
*/ */
@Test @Test
public void testEvictChangedLease() throws LeasingException { public void testEvictChangedLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
KinesisClientLease lease = builder.withLease("1").build().get("1"); Lease lease = builder.withLease("1").build().get("1");
// Change the owner only - this should cause our optimistic lock to fail. // Change the owner only - this should cause our optimistic lock to fail.
lease.setLeaseOwner("otherOwner"); lease.leaseOwner("otherOwner");
Assert.assertFalse(leaseManager.evictLease(lease)); assertFalse(leaseRefresher.evictLease(lease));
} }
/** /**
@ -210,13 +220,13 @@ public class DynamoDBLeaseManagerIntegrationTest extends LeaseIntegrationTest {
*/ */
@Test @Test
public void testDeleteLease() throws LeasingException { public void testDeleteLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
KinesisClientLease lease = builder.withLease("1").build().get("1"); Lease lease = builder.withLease("1").build().get("1");
leaseManager.deleteLease(lease); leaseRefresher.deleteLease(lease);
KinesisClientLease newLease = leaseManager.getLease(lease.getLeaseKey()); Lease newLease = leaseRefresher.getLease(lease.leaseKey());
Assert.assertNull(newLease); assertNull(newLease);
} }
/** /**
@ -224,26 +234,25 @@ public class DynamoDBLeaseManagerIntegrationTest extends LeaseIntegrationTest {
*/ */
@Test @Test
public void testDeleteNonexistentLease() throws LeasingException { public void testDeleteNonexistentLease() throws LeasingException {
KinesisClientLease lease = new KinesisClientLease(); Lease lease = new Lease();
lease.setLeaseKey("1"); lease.leaseKey("1");
// The lease has not been written to DDB - try to delete it and expect success. // The lease has not been written to DDB - try to delete it and expect success.
leaseManager.deleteLease(lease); leaseRefresher.deleteLease(lease);
} }
@Test @Test
public void testWaitUntilLeaseTableExists() throws LeasingException { public void testWaitUntilLeaseTableExists() throws LeasingException {
KinesisClientDynamoDBLeaseManager manager = new KinesisClientDynamoDBLeaseManager("nagl_ShardProgress", ddbClient, true) { DynamoDBLeaseRefresher refresher = new DynamoDBLeaseRefresher("nagl_ShardProgress", ddbClient, new DynamoDBLeaseSerializer(), true) {
@Override @Override
long sleep(long timeToSleepMillis) { long sleep(long timeToSleepMillis) {
Assert.fail("Should not sleep"); fail("Should not sleep");
return 0L; return 0L;
} }
}; };
Assert.assertTrue(manager.waitUntilLeaseTableExists(1, 1)); assertTrue(refresher.waitUntilLeaseTableExists(1, 1));
} }
@Test @Test
@ -252,18 +261,17 @@ public class DynamoDBLeaseManagerIntegrationTest extends LeaseIntegrationTest {
* Just using AtomicInteger for the indirection it provides. * Just using AtomicInteger for the indirection it provides.
*/ */
final AtomicInteger sleepCounter = new AtomicInteger(0); final AtomicInteger sleepCounter = new AtomicInteger(0);
KinesisClientDynamoDBLeaseManager manager = new KinesisClientDynamoDBLeaseManager("nonexistentTable", ddbClient, true) { DynamoDBLeaseRefresher refresher = new DynamoDBLeaseRefresher("nonexistentTable", ddbClient, new DynamoDBLeaseSerializer(), true) {
@Override @Override
long sleep(long timeToSleepMillis) { long sleep(long timeToSleepMillis) {
Assert.assertEquals(1000L, timeToSleepMillis); assertEquals(1000L, timeToSleepMillis);
sleepCounter.incrementAndGet(); sleepCounter.incrementAndGet();
return 1000L; return 1000L;
} }
}; };
Assert.assertFalse(manager.waitUntilLeaseTableExists(2, 1)); assertFalse(refresher.waitUntilLeaseTableExists(2, 1));
Assert.assertEquals(1, sleepCounter.get()); assertEquals(1, sleepCounter.get());
} }
} }

View file

@ -12,34 +12,37 @@
* express or implied. See the License for the specific language governing * express or implied. See the License for the specific language governing
* permissions and limitations under the License. * permissions and limitations under the License.
*/ */
package software.amazon.kinesis.leases; package software.amazon.kinesis.leases.dynamodb;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
import software.amazon.kinesis.leases.exceptions.LeasingException;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import java.util.Collections; import java.util.Collections;
import java.util.Map; import java.util.Map;
import java.util.concurrent.Executors; import java.util.concurrent.Executors;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.LeaseIntegrationTest;
import software.amazon.kinesis.leases.LeaseRenewer;
import software.amazon.kinesis.leases.exceptions.LeasingException;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
public class DynamoDBLeaseRenewerIntegrationTest extends LeaseIntegrationTest { public class DynamoDBLeaseRenewerIntegrationTest extends LeaseIntegrationTest {
// This test case's leases last 2 seconds // This test case's leases last 2 seconds
private static final long LEASE_DURATION_MILLIS = 2000L; private static final long LEASE_DURATION_MILLIS = 2000L;
private LeaseRenewer<KinesisClientLease> renewer; private LeaseRenewer renewer;
@Before @Before
public void setUp() { public void setUp() {
renewer = new DynamoDBLeaseRenewer<KinesisClientLease>( renewer = new DynamoDBLeaseRenewer(leaseRefresher, "foo", LEASE_DURATION_MILLIS, Executors.newCachedThreadPool());
leaseManager, "foo", LEASE_DURATION_MILLIS, Executors.newCachedThreadPool());
} }
@Test @Test
public void testSimpleRenew() throws LeasingException { public void testSimpleRenew() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", "foo").build(); builder.withLease("1", "foo").build();
@ -49,22 +52,22 @@ public class DynamoDBLeaseRenewerIntegrationTest extends LeaseIntegrationTest {
@Test @Test
public void testLeaseLoss() throws LeasingException { public void testLeaseLoss() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", "foo").withLease("2", "foo").build(); builder.withLease("1", "foo").withLease("2", "foo").build();
builder.addLeasesToRenew(renewer, "1", "2"); builder.addLeasesToRenew(renewer, "1", "2");
KinesisClientLease renewedLease = builder.renewMutateAssert(renewer, "1", "2").get("2"); Lease renewedLease = builder.renewMutateAssert(renewer, "1", "2").get("2");
// lose lease 2 // lose lease 2
leaseManager.takeLease(renewedLease, "bar"); leaseRefresher.takeLease(renewedLease, "bar");
builder.renewMutateAssert(renewer, "1"); builder.renewMutateAssert(renewer, "1");
} }
@Test @Test
public void testClear() throws LeasingException { public void testClear() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", "foo").build(); builder.withLease("1", "foo").build();
builder.addLeasesToRenew(renewer, "1"); builder.addLeasesToRenew(renewer, "1");
@ -76,143 +79,143 @@ public class DynamoDBLeaseRenewerIntegrationTest extends LeaseIntegrationTest {
@Test @Test
public void testGetCurrentlyHeldLease() throws LeasingException { public void testGetCurrentlyHeldLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", "foo").build(); builder.withLease("1", "foo").build();
builder.addLeasesToRenew(renewer, "1"); builder.addLeasesToRenew(renewer, "1");
builder.renewMutateAssert(renewer, "1"); builder.renewMutateAssert(renewer, "1");
// this should be a copy that doesn't get updated // this should be a copy that doesn't get updated
KinesisClientLease lease = renewer.getCurrentlyHeldLease("1"); Lease lease = renewer.getCurrentlyHeldLease("1");
Assert.assertEquals((Long) 1L, lease.getLeaseCounter()); Assert.assertEquals((Long) 1L, lease.leaseCounter());
// do one renewal and make sure the old copy doesn't get updated // do one renewal and make sure the old copy doesn't get updated
builder.renewMutateAssert(renewer, "1"); builder.renewMutateAssert(renewer, "1");
Assert.assertEquals((Long) 1L, lease.getLeaseCounter()); Assert.assertEquals((Long) 1L, lease.leaseCounter());
} }
@Test @Test
public void testGetCurrentlyHeldLeases() throws LeasingException { public void testGetCurrentlyHeldLeases() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", "foo").withLease("2", "foo").build(); builder.withLease("1", "foo").withLease("2", "foo").build();
builder.addLeasesToRenew(renewer, "1", "2"); builder.addLeasesToRenew(renewer, "1", "2");
KinesisClientLease lease2 = builder.renewMutateAssert(renewer, "1", "2").get("2"); Lease lease2 = builder.renewMutateAssert(renewer, "1", "2").get("2");
// This should be a copy that doesn't get updated // This should be a copy that doesn't get updated
Map<String, KinesisClientLease> heldLeases = renewer.getCurrentlyHeldLeases(); Map<String, Lease> heldLeases = renewer.getCurrentlyHeldLeases();
Assert.assertEquals(2, heldLeases.size()); Assert.assertEquals(2, heldLeases.size());
Assert.assertEquals((Long) 1L, heldLeases.get("1").getLeaseCounter()); Assert.assertEquals((Long) 1L, heldLeases.get("1").leaseCounter());
Assert.assertEquals((Long) 1L, heldLeases.get("2").getLeaseCounter()); Assert.assertEquals((Long) 1L, heldLeases.get("2").leaseCounter());
// lose lease 2 // lose lease 2
leaseManager.takeLease(lease2, "bar"); leaseRefresher.takeLease(lease2, "bar");
// Do another renewal and make sure the copy doesn't change // Do another renewal and make sure the copy doesn't change
builder.renewMutateAssert(renewer, "1"); builder.renewMutateAssert(renewer, "1");
Assert.assertEquals(2, heldLeases.size()); Assert.assertEquals(2, heldLeases.size());
Assert.assertEquals((Long) 1L, heldLeases.get("1").getLeaseCounter()); Assert.assertEquals((Long) 1L, heldLeases.get("1").leaseCounter());
Assert.assertEquals((Long) 1L, heldLeases.get("2").getLeaseCounter()); Assert.assertEquals((Long) 1L, heldLeases.get("2").leaseCounter());
} }
@Test @Test
public void testUpdateLease() throws LeasingException { public void testUpdateLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", "foo").build(); builder.withLease("1", "foo").build();
builder.addLeasesToRenew(renewer, "1"); builder.addLeasesToRenew(renewer, "1");
builder.renewMutateAssert(renewer, "1"); builder.renewMutateAssert(renewer, "1");
KinesisClientLease expected = renewer.getCurrentlyHeldLease("1"); Lease expected = renewer.getCurrentlyHeldLease("1");
expected.setCheckpoint(new ExtendedSequenceNumber("new checkpoint")); expected.checkpoint(new ExtendedSequenceNumber("new checkpoint"));
Assert.assertTrue(renewer.updateLease(expected, expected.getConcurrencyToken())); Assert.assertTrue(renewer.updateLease(expected, expected.concurrencyToken()));
// Assert that the counter and data have changed immediately after the update... // Assert that the counter and data have changed immediately after the update...
KinesisClientLease actual = renewer.getCurrentlyHeldLease("1"); Lease actual = renewer.getCurrentlyHeldLease("1");
expected.setLeaseCounter(expected.getLeaseCounter() + 1); expected.leaseCounter(expected.leaseCounter() + 1);
Assert.assertEquals(expected, actual); Assert.assertEquals(expected, actual);
// ...and after another round of renewal // ...and after another round of renewal
renewer.renewLeases(); renewer.renewLeases();
actual = renewer.getCurrentlyHeldLease("1"); actual = renewer.getCurrentlyHeldLease("1");
expected.setLeaseCounter(expected.getLeaseCounter() + 1); expected.leaseCounter(expected.leaseCounter() + 1);
Assert.assertEquals(expected, actual); Assert.assertEquals(expected, actual);
} }
@Test @Test
public void testUpdateLostLease() throws LeasingException { public void testUpdateLostLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", "foo").build(); builder.withLease("1", "foo").build();
builder.addLeasesToRenew(renewer, "1"); builder.addLeasesToRenew(renewer, "1");
builder.renewMutateAssert(renewer, "1"); builder.renewMutateAssert(renewer, "1");
KinesisClientLease lease = renewer.getCurrentlyHeldLease("1"); Lease lease = renewer.getCurrentlyHeldLease("1");
// cause lease loss such that the renewer doesn't realize he's lost the lease when update is called // cause lease loss such that the renewer doesn't realize he's lost the lease when update is called
leaseManager.renewLease(lease); leaseRefresher.renewLease(lease);
// renewer still thinks he has the lease // renewer still thinks he has the lease
Assert.assertNotNull(renewer.getCurrentlyHeldLease("1")); Assert.assertNotNull(renewer.getCurrentlyHeldLease("1"));
lease.setCheckpoint(new ExtendedSequenceNumber("new checkpoint")); lease.checkpoint(new ExtendedSequenceNumber("new checkpoint"));
// update fails // update fails
Assert.assertFalse(renewer.updateLease(lease, lease.getConcurrencyToken())); Assert.assertFalse(renewer.updateLease(lease, lease.concurrencyToken()));
// renewer no longer thinks he has the lease // renewer no longer thinks he has the lease
Assert.assertNull(renewer.getCurrentlyHeldLease("1")); Assert.assertNull(renewer.getCurrentlyHeldLease("1"));
} }
@Test @Test
public void testUpdateOldLease() throws LeasingException { public void testUpdateOldLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", "foo").build(); builder.withLease("1", "foo").build();
builder.addLeasesToRenew(renewer, "1"); builder.addLeasesToRenew(renewer, "1");
builder.renewMutateAssert(renewer, "1"); builder.renewMutateAssert(renewer, "1");
KinesisClientLease lease = renewer.getCurrentlyHeldLease("1"); Lease lease = renewer.getCurrentlyHeldLease("1");
// cause lease loss such that the renewer knows the lease has been lost when update is called // cause lease loss such that the renewer knows the lease has been lost when update is called
leaseManager.takeLease(lease, "bar"); leaseRefresher.takeLease(lease, "bar");
builder.renewMutateAssert(renewer); builder.renewMutateAssert(renewer);
lease.setCheckpoint(new ExtendedSequenceNumber("new checkpoint")); lease.checkpoint(new ExtendedSequenceNumber("new checkpoint"));
Assert.assertFalse(renewer.updateLease(lease, lease.getConcurrencyToken())); Assert.assertFalse(renewer.updateLease(lease, lease.concurrencyToken()));
} }
@Test @Test
public void testUpdateRegainedLease() throws LeasingException { public void testUpdateRegainedLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", "foo").build(); builder.withLease("1", "foo").build();
builder.addLeasesToRenew(renewer, "1"); builder.addLeasesToRenew(renewer, "1");
builder.renewMutateAssert(renewer, "1"); builder.renewMutateAssert(renewer, "1");
KinesisClientLease lease = renewer.getCurrentlyHeldLease("1"); Lease lease = renewer.getCurrentlyHeldLease("1");
// cause lease loss such that the renewer knows the lease has been lost when update is called // cause lease loss such that the renewer knows the lease has been lost when update is called
leaseManager.takeLease(lease, "bar"); leaseRefresher.takeLease(lease, "bar");
builder.renewMutateAssert(renewer); builder.renewMutateAssert(renewer);
// regain the lease // regain the lease
builder.addLeasesToRenew(renewer, "1"); builder.addLeasesToRenew(renewer, "1");
lease.setCheckpoint(new ExtendedSequenceNumber("new checkpoint")); lease.checkpoint(new ExtendedSequenceNumber("new checkpoint"));
Assert.assertFalse(renewer.updateLease(lease, lease.getConcurrencyToken())); Assert.assertFalse(renewer.updateLease(lease, lease.concurrencyToken()));
} }
@Test @Test
public void testIgnoreNoRenewalTimestamp() throws LeasingException { public void testIgnoreNoRenewalTimestamp() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
KinesisClientLease lease = builder.withLease("1", "foo").build().get("1"); Lease lease = builder.withLease("1", "foo").build().get("1");
lease.setLastCounterIncrementNanos(null); lease.lastCounterIncrementNanos(null);
renewer.addLeasesToRenew(Collections.singleton(lease)); renewer.addLeasesToRenew(Collections.singleton(lease));
@ -221,7 +224,7 @@ public class DynamoDBLeaseRenewerIntegrationTest extends LeaseIntegrationTest {
@Test @Test
public void testLeaseTimeout() throws LeasingException, InterruptedException { public void testLeaseTimeout() throws LeasingException, InterruptedException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", "foo").build(); builder.withLease("1", "foo").build();
@ -239,13 +242,13 @@ public class DynamoDBLeaseRenewerIntegrationTest extends LeaseIntegrationTest {
final String shardId = "shd-0-0"; final String shardId = "shd-0-0";
final String owner = "foo:8000"; final String owner = "foo:8000";
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease(shardId, owner); builder.withLease(shardId, owner);
Map<String, KinesisClientLease> leases = builder.build(); Map<String, Lease> leases = builder.build();
DynamoDBLeaseRenewer<KinesisClientLease> renewer =new DynamoDBLeaseRenewer<KinesisClientLease>( DynamoDBLeaseRenewer renewer = new DynamoDBLeaseRenewer(
leaseManager, owner, 30000L, Executors.newCachedThreadPool()); leaseRefresher, owner, 30000L, Executors.newCachedThreadPool());
renewer.initialize(); renewer.initialize();
Map<String, KinesisClientLease> heldLeases = renewer.getCurrentlyHeldLeases(); Map<String, Lease> heldLeases = renewer.getCurrentlyHeldLeases();
Assert.assertEquals(leases.size(), heldLeases.size()); Assert.assertEquals(leases.size(), heldLeases.size());
Assert.assertEquals(leases.keySet(), heldLeases.keySet()); Assert.assertEquals(leases.keySet(), heldLeases.keySet());
} }

View file

@ -12,61 +12,56 @@
* express or implied. See the License for the specific language governing * express or implied. See the License for the specific language governing
* permissions and limitations under the License. * permissions and limitations under the License.
*/ */
package software.amazon.kinesis.leases; package software.amazon.kinesis.leases.dynamodb;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.Arrays; import java.util.Arrays;
import java.util.HashSet;
import java.util.List; import java.util.List;
import java.util.UUID; import java.util.UUID;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors; import java.util.concurrent.Executors;
import org.junit.After; import org.junit.After;
import org.junit.Assert;
import org.junit.Before; import org.junit.Before;
import org.junit.Test; import org.junit.Test;
import org.mockito.Mockito; import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.runners.MockitoJUnitRunner;
import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseRenewer;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException; import software.amazon.kinesis.leases.exceptions.ProvisionedThroughputException;
@RunWith(MockitoJUnitRunner.class)
public class DynamoDBLeaseRenewerTest { public class DynamoDBLeaseRenewerTest {
private final String workerIdentifier = "WorkerId";
private final long leaseDurationMillis = 10000;
private DynamoDBLeaseRenewer renewer;
private List<Lease> leasesToRenew;
LeaseManager<Lease> leaseManager; @Mock
String workerIdentifier; private LeaseRefresher leaseRefresher;
long leaseDurationMillis;
ExecutorService leaseRenewalExecService;
DynamoDBLeaseRenewer<Lease> renewer;
List<Lease> leasesToRenew;
private static Lease newLease(String leaseKey,
String leaseOwner,
Long leaseCounter,
UUID concurrencyToken,
Long lastCounterIncrementNanos) {
Lease lease = new Lease();
lease.setLeaseKey(leaseKey);
lease.setLeaseOwner(leaseOwner);
lease.setLeaseCounter(leaseCounter);
lease.setConcurrencyToken(concurrencyToken);
lease.setLastCounterIncrementNanos(lastCounterIncrementNanos);
return lease;
}
private static Lease newLease(String leaseKey) { private static Lease newLease(String leaseKey) {
return newLease(leaseKey, "leaseOwner", 0L, UUID.randomUUID(), System.nanoTime()); return new Lease(leaseKey, "LeaseOwner", 0L, UUID.randomUUID(), System.nanoTime(), null, null, null, new HashSet<>());
} }
@SuppressWarnings("unchecked")
@Before @Before
public void before() { public void before() {
leaseManager = Mockito.mock(LeaseManager.class);
workerIdentifier = "workerId";
leaseDurationMillis = 10000;
leaseRenewalExecService = Executors.newSingleThreadExecutor();
leasesToRenew = null; leasesToRenew = null;
renewer = new DynamoDBLeaseRenewer<>(leaseManager, renewer = new DynamoDBLeaseRenewer(leaseRefresher,
workerIdentifier, workerIdentifier,
leaseDurationMillis, leaseDurationMillis,
Executors.newCachedThreadPool()); Executors.newCachedThreadPool());
@ -77,8 +72,8 @@ public class DynamoDBLeaseRenewerTest {
if (leasesToRenew == null) { if (leasesToRenew == null) {
return; return;
} }
for (Lease l : leasesToRenew) { for (Lease lease : leasesToRenew) {
Mockito.verify(leaseManager, Mockito.times(1)).renewLease(l); verify(leaseRefresher, times(1)).renewLease(eq(lease));
} }
} }
@ -91,16 +86,15 @@ public class DynamoDBLeaseRenewerTest {
*/ */
Lease lease1 = newLease("1"); Lease lease1 = newLease("1");
Lease lease2 = newLease("2"); Lease lease2 = newLease("2");
leasesToRenew = leasesToRenew = Arrays.asList(lease1,lease2);
Arrays.asList(lease1,lease2);
renewer.addLeasesToRenew(leasesToRenew); renewer.addLeasesToRenew(leasesToRenew);
Mockito.doReturn(true).when(leaseManager).renewLease(lease1); doReturn(true).when(leaseRefresher).renewLease(lease1);
Mockito.doReturn(true).when(leaseManager).renewLease(lease2); doReturn(true).when(leaseRefresher).renewLease(lease2);
renewer.renewLeases(); renewer.renewLeases();
Assert.assertEquals(2, renewer.getCurrentlyHeldLeases().size()); assertEquals(2, renewer.getCurrentlyHeldLeases().size());
} }
@Test @Test
@ -108,19 +102,19 @@ public class DynamoDBLeaseRenewerTest {
String leaseKey = "expiredLease"; String leaseKey = "expiredLease";
long initialCounterIncrementNanos = 5L; // "expired" time. long initialCounterIncrementNanos = 5L; // "expired" time.
Lease lease1 = newLease(leaseKey); Lease lease1 = newLease(leaseKey);
lease1.setLastCounterIncrementNanos(initialCounterIncrementNanos); lease1.lastCounterIncrementNanos(initialCounterIncrementNanos);
leasesToRenew = new ArrayList<>(); leasesToRenew = new ArrayList<>();
leasesToRenew.add(lease1); leasesToRenew.add(lease1);
Mockito.doReturn(true).when(leaseManager).renewLease(lease1); doReturn(true).when(leaseRefresher).renewLease(lease1);
renewer.addLeasesToRenew(leasesToRenew); renewer.addLeasesToRenew(leasesToRenew);
Assert.assertTrue(lease1.isExpired(1, System.nanoTime())); assertTrue(lease1.isExpired(1, System.nanoTime()));
Assert.assertNull(renewer.getCurrentlyHeldLease(leaseKey)); assertNull(renewer.getCurrentlyHeldLease(leaseKey));
renewer.renewLeases(); renewer.renewLeases();
// Don't renew lease(s) with same key if getCurrentlyHeldLease returned null previously // Don't renew lease(s) with same key if getCurrentlyHeldLease returned null previously
Assert.assertNull(renewer.getCurrentlyHeldLease(leaseKey)); assertNull(renewer.getCurrentlyHeldLease(leaseKey));
Assert.assertFalse(renewer.getCurrentlyHeldLeases().containsKey(leaseKey)); assertFalse(renewer.getCurrentlyHeldLeases().containsKey(leaseKey));
// Clear the list to avoid triggering expectation mismatch in after(). // Clear the list to avoid triggering expectation mismatch in after().
leasesToRenew.clear(); leasesToRenew.clear();

View file

@ -12,7 +12,7 @@
* express or implied. See the License for the specific language governing * express or implied. See the License for the specific language governing
* permissions and limitations under the License. * permissions and limitations under the License.
*/ */
package software.amazon.kinesis.leases; package software.amazon.kinesis.leases.dynamodb;
import java.util.Map; import java.util.Map;
@ -20,21 +20,23 @@ import org.junit.Assert;
import org.junit.Before; import org.junit.Before;
import org.junit.Test; import org.junit.Test;
import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.LeaseIntegrationTest;
import software.amazon.kinesis.leases.exceptions.LeasingException; import software.amazon.kinesis.leases.exceptions.LeasingException;
public class DynamoDBLeaseTakerIntegrationTest extends LeaseIntegrationTest { public class DynamoDBLeaseTakerIntegrationTest extends LeaseIntegrationTest {
private static final long LEASE_DURATION_MILLIS = 1000L; private static final long LEASE_DURATION_MILLIS = 1000L;
private DynamoDBLeaseTaker<KinesisClientLease> taker; private DynamoDBLeaseTaker taker;
@Before @Before
public void setUp() { public void setUp() {
taker = new DynamoDBLeaseTaker<KinesisClientLease>(leaseManager, "foo", LEASE_DURATION_MILLIS); taker = new DynamoDBLeaseTaker(leaseRefresher, "foo", LEASE_DURATION_MILLIS);
} }
@Test @Test
public void testSimpleLeaseTake() throws LeasingException { public void testSimpleLeaseTake() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", null).build(); builder.withLease("1", null).build();
@ -43,7 +45,7 @@ public class DynamoDBLeaseTakerIntegrationTest extends LeaseIntegrationTest {
@Test @Test
public void testNotTakeUpdatedLease() throws LeasingException { public void testNotTakeUpdatedLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", "bar").build(); builder.withLease("1", "bar").build();
@ -56,7 +58,7 @@ public class DynamoDBLeaseTakerIntegrationTest extends LeaseIntegrationTest {
@Test @Test
public void testTakeOwnLease() throws LeasingException { public void testTakeOwnLease() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", taker.getWorkerIdentifier()).build(); builder.withLease("1", taker.getWorkerIdentifier()).build();
@ -67,7 +69,7 @@ public class DynamoDBLeaseTakerIntegrationTest extends LeaseIntegrationTest {
@Test @Test
public void testNotTakeNewOwnedLease() throws LeasingException, InterruptedException { public void testNotTakeNewOwnedLease() throws LeasingException, InterruptedException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", "bar").build(); builder.withLease("1", "bar").build();
@ -85,7 +87,7 @@ public class DynamoDBLeaseTakerIntegrationTest extends LeaseIntegrationTest {
*/ */
@Test @Test
public void testNonGreedyTake() throws LeasingException { public void testNonGreedyTake() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
for (int i = 0; i < 3; i++) { for (int i = 0; i < 3; i++) {
builder.withLease(Integer.toString(i), null); builder.withLease(Integer.toString(i), null);
@ -103,7 +105,7 @@ public class DynamoDBLeaseTakerIntegrationTest extends LeaseIntegrationTest {
*/ */
@Test @Test
public void testNoStealWhenOffByOne() throws LeasingException { public void testNoStealWhenOffByOne() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", "bar") builder.withLease("1", "bar")
.withLease("2", "bar") .withLease("2", "bar")
@ -124,7 +126,7 @@ public class DynamoDBLeaseTakerIntegrationTest extends LeaseIntegrationTest {
*/ */
@Test @Test
public void testSteal() throws LeasingException { public void testSteal() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", "bar"); builder.withLease("1", "bar");
for (int i = 2; i <= 6; i++) { for (int i = 2; i <= 6; i++) {
@ -135,7 +137,7 @@ public class DynamoDBLeaseTakerIntegrationTest extends LeaseIntegrationTest {
builder.build(); builder.build();
// Assert that one lease was stolen from baz. // Assert that one lease was stolen from baz.
Map<String, KinesisClientLease> takenLeases = builder.takeMutateAssert(taker, 1); Map<String, Lease> takenLeases = builder.takeMutateAssert(taker, 1);
// Assert that it was one of baz's leases (shardId != 1) // Assert that it was one of baz's leases (shardId != 1)
String shardIdStolen = takenLeases.keySet().iterator().next(); String shardIdStolen = takenLeases.keySet().iterator().next();
@ -148,7 +150,7 @@ public class DynamoDBLeaseTakerIntegrationTest extends LeaseIntegrationTest {
*/ */
@Test @Test
public void testNoStealWhenExpiredLeases() throws LeasingException { public void testNoStealWhenExpiredLeases() throws LeasingException {
TestHarnessBuilder builder = new TestHarnessBuilder(leaseManager); TestHarnessBuilder builder = new TestHarnessBuilder(leaseRefresher);
builder.withLease("1", null); builder.withLease("1", null);
for (int i = 2; i <= 4; i++) { for (int i = 2; i <= 4; i++) {

View file

@ -12,7 +12,7 @@
* express or implied. See the License for the specific language governing * express or implied. See the License for the specific language governing
* permissions and limitations under the License. * permissions and limitations under the License.
*/ */
package software.amazon.kinesis.leases; package software.amazon.kinesis.leases.dynamodb;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
@ -24,6 +24,7 @@ import org.junit.AfterClass;
import org.junit.Before; import org.junit.Before;
import org.junit.BeforeClass; import org.junit.BeforeClass;
import org.junit.Test; import org.junit.Test;
import software.amazon.kinesis.leases.dynamodb.DynamoDBLeaseTaker;
/** /**
* *

View file

@ -0,0 +1,179 @@
/*
* Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Amazon Software License (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/asl/
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*/
package software.amazon.kinesis.leases.dynamodb;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.LeaseRenewer;
import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException;
import software.amazon.kinesis.leases.exceptions.LeasingException;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
public class TestHarnessBuilder {
private long currentTimeNanos;
private Map<String, Lease> leases = new HashMap<>();
private DynamoDBLeaseRefresher leaseRefresher;
private Map<String, Lease> originalLeases = new HashMap<>();
private Callable<Long> timeProvider = new Callable<Long>() {
@Override
public Long call() throws Exception {
return currentTimeNanos;
}
};
public TestHarnessBuilder(final DynamoDBLeaseRefresher leaseRefresher) {
this.leaseRefresher = leaseRefresher;
}
public TestHarnessBuilder withLease(String shardId) {
return withLease(shardId, "leaseOwner");
}
public TestHarnessBuilder withLease(String shardId, String owner) {
Lease lease = createLease(shardId, owner);
Lease originalLease = createLease(shardId, owner);
leases.put(shardId, lease);
originalLeases.put(shardId, originalLease);
return this;
}
private Lease createLease(String shardId, String owner) {
Lease lease = new Lease();
lease.checkpoint(new ExtendedSequenceNumber("checkpoint"));
lease.ownerSwitchesSinceCheckpoint(0L);
lease.leaseCounter(0L);
lease.leaseOwner(owner);
lease.parentShardIds(Collections.singleton("parentShardId"));
lease.leaseKey(shardId);
return lease;
}
public Map<String, Lease> build() throws LeasingException {
for (Lease lease : leases.values()) {
leaseRefresher.createLeaseIfNotExists(lease);
if (lease.leaseOwner() != null) {
lease.lastCounterIncrementNanos(System.nanoTime());
}
}
currentTimeNanos = System.nanoTime();
return leases;
}
public void passTime(long millis) {
currentTimeNanos += millis * 1000000;
}
public Map<String, Lease> takeMutateAssert(DynamoDBLeaseTaker taker, int numToTake)
throws LeasingException {
Map<String, Lease> result = taker.takeLeases(timeProvider);
assertEquals(numToTake, result.size());
for (Lease actual : result.values()) {
Lease original = leases.get(actual.leaseKey());
assertNotNull(original);
mutateAssert(taker.getWorkerIdentifier(), original, actual);
}
return result;
}
public Map<String, Lease> takeMutateAssert(DynamoDBLeaseTaker taker, String... takenShardIds)
throws LeasingException {
Map<String, Lease> result = taker.takeLeases(timeProvider);
assertEquals(takenShardIds.length, result.size());
for (String shardId : takenShardIds) {
Lease original = leases.get(shardId);
assertNotNull(original);
Lease actual = result.get(shardId);
assertNotNull(actual);
mutateAssert(taker.getWorkerIdentifier(), original, actual);
}
return result;
}
private void mutateAssert(String newWorkerIdentifier, Lease original, Lease actual) {
original.leaseCounter(original.leaseCounter() + 1);
if (original.leaseOwner() != null && !newWorkerIdentifier.equals(original.leaseOwner())) {
original.ownerSwitchesSinceCheckpoint(original.ownerSwitchesSinceCheckpoint() + 1);
}
original.leaseOwner(newWorkerIdentifier);
assertEquals(original, actual); // Assert the contents of the lease
}
public void addLeasesToRenew(LeaseRenewer renewer, String... shardIds)
throws DependencyException, InvalidStateException {
List<Lease> leasesToRenew = new ArrayList<Lease>();
for (String shardId : shardIds) {
Lease lease = leases.get(shardId);
assertNotNull(lease);
leasesToRenew.add(lease);
}
renewer.addLeasesToRenew(leasesToRenew);
}
public Map<String, Lease> renewMutateAssert(LeaseRenewer renewer, String... renewedShardIds)
throws DependencyException, InvalidStateException {
renewer.renewLeases();
Map<String, Lease> heldLeases = renewer.getCurrentlyHeldLeases();
assertEquals(renewedShardIds.length, heldLeases.size());
for (String shardId : renewedShardIds) {
Lease original = originalLeases.get(shardId);
assertNotNull(original);
Lease actual = heldLeases.get(shardId);
assertNotNull(actual);
original.leaseCounter(original.leaseCounter() + 1);
assertEquals(original, actual);
}
return heldLeases;
}
public void renewAllLeases() throws LeasingException {
for (Lease lease : leases.values()) {
leaseRefresher.renewLease(lease);
}
}
}

View file

@ -26,8 +26,8 @@ import java.util.List;
import org.junit.Before; import org.junit.Before;
import org.junit.Test; import org.junit.Test;
import software.amazon.kinesis.leases.LeaseManager; import software.amazon.kinesis.leases.Lease;
import software.amazon.kinesis.leases.KinesisClientLease; import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.ShardInfo; import software.amazon.kinesis.leases.ShardInfo;
import software.amazon.kinesis.leases.exceptions.DependencyException; import software.amazon.kinesis.leases.exceptions.DependencyException;
import software.amazon.kinesis.leases.exceptions.InvalidStateException; import software.amazon.kinesis.leases.exceptions.InvalidStateException;
@ -58,10 +58,10 @@ public class BlockOnParentShardTaskTest {
@Test @Test
public final void testCallNoParents() public final void testCallNoParents()
throws DependencyException, InvalidStateException, ProvisionedThroughputException { throws DependencyException, InvalidStateException, ProvisionedThroughputException {
LeaseManager<KinesisClientLease> leaseManager = mock(LeaseManager.class); LeaseRefresher leaseRefresher = mock(LeaseRefresher.class);
when(leaseManager.getLease(shardId)).thenReturn(null); when(leaseRefresher.getLease(shardId)).thenReturn(null);
BlockOnParentShardTask task = new BlockOnParentShardTask(shardInfo, leaseManager, backoffTimeInMillis); BlockOnParentShardTask task = new BlockOnParentShardTask(shardInfo, leaseRefresher, backoffTimeInMillis);
TaskResult result = task.call(); TaskResult result = task.call();
assertNull(result.getException()); assertNull(result.getException());
} }
@ -83,26 +83,26 @@ public class BlockOnParentShardTaskTest {
List<String> parentShardIds = new ArrayList<>(); List<String> parentShardIds = new ArrayList<>();
TaskResult result = null; TaskResult result = null;
KinesisClientLease parent1Lease = new KinesisClientLease(); Lease parent1Lease = new Lease();
parent1Lease.setCheckpoint(ExtendedSequenceNumber.SHARD_END); parent1Lease.checkpoint(ExtendedSequenceNumber.SHARD_END);
KinesisClientLease parent2Lease = new KinesisClientLease(); Lease parent2Lease = new Lease();
parent2Lease.setCheckpoint(ExtendedSequenceNumber.SHARD_END); parent2Lease.checkpoint(ExtendedSequenceNumber.SHARD_END);
LeaseManager<KinesisClientLease> leaseManager = mock(LeaseManager.class); LeaseRefresher leaseRefresher = mock(LeaseRefresher.class);
when(leaseManager.getLease(parent1ShardId)).thenReturn(parent1Lease); when(leaseRefresher.getLease(parent1ShardId)).thenReturn(parent1Lease);
when(leaseManager.getLease(parent2ShardId)).thenReturn(parent2Lease); when(leaseRefresher.getLease(parent2ShardId)).thenReturn(parent2Lease);
// test single parent // test single parent
parentShardIds.add(parent1ShardId); parentShardIds.add(parent1ShardId);
shardInfo = new ShardInfo(shardId, concurrencyToken, parentShardIds, ExtendedSequenceNumber.TRIM_HORIZON); shardInfo = new ShardInfo(shardId, concurrencyToken, parentShardIds, ExtendedSequenceNumber.TRIM_HORIZON);
task = new BlockOnParentShardTask(shardInfo, leaseManager, backoffTimeInMillis); task = new BlockOnParentShardTask(shardInfo, leaseRefresher, backoffTimeInMillis);
result = task.call(); result = task.call();
assertNull(result.getException()); assertNull(result.getException());
// test two parents // test two parents
parentShardIds.add(parent2ShardId); parentShardIds.add(parent2ShardId);
shardInfo = new ShardInfo(shardId, concurrencyToken, parentShardIds, ExtendedSequenceNumber.TRIM_HORIZON); shardInfo = new ShardInfo(shardId, concurrencyToken, parentShardIds, ExtendedSequenceNumber.TRIM_HORIZON);
task = new BlockOnParentShardTask(shardInfo, leaseManager, backoffTimeInMillis); task = new BlockOnParentShardTask(shardInfo, leaseRefresher, backoffTimeInMillis);
result = task.call(); result = task.call();
assertNull(result.getException()); assertNull(result.getException());
} }
@ -124,27 +124,27 @@ public class BlockOnParentShardTaskTest {
List<String> parentShardIds = new ArrayList<>(); List<String> parentShardIds = new ArrayList<>();
TaskResult result = null; TaskResult result = null;
KinesisClientLease parent1Lease = new KinesisClientLease(); Lease parent1Lease = new Lease();
parent1Lease.setCheckpoint(ExtendedSequenceNumber.LATEST); parent1Lease.checkpoint(ExtendedSequenceNumber.LATEST);
KinesisClientLease parent2Lease = new KinesisClientLease(); Lease parent2Lease = new Lease();
// mock a sequence number checkpoint // mock a sequence number checkpoint
parent2Lease.setCheckpoint(new ExtendedSequenceNumber("98182584034")); parent2Lease.checkpoint(new ExtendedSequenceNumber("98182584034"));
LeaseManager<KinesisClientLease> leaseManager = mock(LeaseManager.class); LeaseRefresher leaseRefresher = mock(LeaseRefresher.class);
when(leaseManager.getLease(parent1ShardId)).thenReturn(parent1Lease); when(leaseRefresher.getLease(parent1ShardId)).thenReturn(parent1Lease);
when(leaseManager.getLease(parent2ShardId)).thenReturn(parent2Lease); when(leaseRefresher.getLease(parent2ShardId)).thenReturn(parent2Lease);
// test single parent // test single parent
parentShardIds.add(parent1ShardId); parentShardIds.add(parent1ShardId);
shardInfo = new ShardInfo(shardId, concurrencyToken, parentShardIds, ExtendedSequenceNumber.TRIM_HORIZON); shardInfo = new ShardInfo(shardId, concurrencyToken, parentShardIds, ExtendedSequenceNumber.TRIM_HORIZON);
task = new BlockOnParentShardTask(shardInfo, leaseManager, backoffTimeInMillis); task = new BlockOnParentShardTask(shardInfo, leaseRefresher, backoffTimeInMillis);
result = task.call(); result = task.call();
assertNotNull(result.getException()); assertNotNull(result.getException());
// test two parents // test two parents
parentShardIds.add(parent2ShardId); parentShardIds.add(parent2ShardId);
shardInfo = new ShardInfo(shardId, concurrencyToken, parentShardIds, ExtendedSequenceNumber.TRIM_HORIZON); shardInfo = new ShardInfo(shardId, concurrencyToken, parentShardIds, ExtendedSequenceNumber.TRIM_HORIZON);
task = new BlockOnParentShardTask(shardInfo, leaseManager, backoffTimeInMillis); task = new BlockOnParentShardTask(shardInfo, leaseRefresher, backoffTimeInMillis);
result = task.call(); result = task.call();
assertNotNull(result.getException()); assertNotNull(result.getException());
} }
@ -165,19 +165,19 @@ public class BlockOnParentShardTaskTest {
parentShardIds.add(parentShardId); parentShardIds.add(parentShardId);
ShardInfo shardInfo = new ShardInfo(shardId, concurrencyToken, parentShardIds, ExtendedSequenceNumber.TRIM_HORIZON); ShardInfo shardInfo = new ShardInfo(shardId, concurrencyToken, parentShardIds, ExtendedSequenceNumber.TRIM_HORIZON);
TaskResult result = null; TaskResult result = null;
KinesisClientLease parentLease = new KinesisClientLease(); Lease parentLease = new Lease();
LeaseManager<KinesisClientLease> leaseManager = mock(LeaseManager.class); LeaseRefresher leaseRefresher = mock(LeaseRefresher.class);
when(leaseManager.getLease(parentShardId)).thenReturn(parentLease); when(leaseRefresher.getLease(parentShardId)).thenReturn(parentLease);
// test when parent shard has not yet been fully processed // test when parent shard has not yet been fully processed
parentLease.setCheckpoint(new ExtendedSequenceNumber("98182584034")); parentLease.checkpoint(new ExtendedSequenceNumber("98182584034"));
task = new BlockOnParentShardTask(shardInfo, leaseManager, backoffTimeInMillis); task = new BlockOnParentShardTask(shardInfo, leaseRefresher, backoffTimeInMillis);
result = task.call(); result = task.call();
assertNotNull(result.getException()); assertNotNull(result.getException());
// test when parent has been fully processed // test when parent has been fully processed
parentLease.setCheckpoint(ExtendedSequenceNumber.SHARD_END); parentLease.checkpoint(ExtendedSequenceNumber.SHARD_END);
task = new BlockOnParentShardTask(shardInfo, leaseManager, backoffTimeInMillis); task = new BlockOnParentShardTask(shardInfo, leaseRefresher, backoffTimeInMillis);
result = task.call(); result = task.call();
assertNull(result.getException()); assertNull(result.getException());
} }

View file

@ -45,14 +45,13 @@ import com.amazonaws.services.kinesis.clientlibrary.lib.worker.InitialPositionIn
import com.amazonaws.services.kinesis.clientlibrary.lib.worker.InitialPositionInStreamExtended; import com.amazonaws.services.kinesis.clientlibrary.lib.worker.InitialPositionInStreamExtended;
import software.amazon.kinesis.checkpoint.RecordProcessorCheckpointer; import software.amazon.kinesis.checkpoint.RecordProcessorCheckpointer;
import software.amazon.kinesis.leases.LeaseManager; import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.KinesisClientLease; import software.amazon.kinesis.leases.ShardDetector;
import software.amazon.kinesis.leases.LeaseManagerProxy;
import software.amazon.kinesis.leases.ShardInfo; import software.amazon.kinesis.leases.ShardInfo;
import software.amazon.kinesis.metrics.IMetricsFactory; import software.amazon.kinesis.metrics.IMetricsFactory;
import software.amazon.kinesis.processor.Checkpointer; import software.amazon.kinesis.processor.Checkpointer;
import software.amazon.kinesis.processor.RecordProcessor;
import software.amazon.kinesis.processor.IRecordProcessorCheckpointer; import software.amazon.kinesis.processor.IRecordProcessorCheckpointer;
import software.amazon.kinesis.processor.RecordProcessor;
import software.amazon.kinesis.retrieval.GetRecordsCache; import software.amazon.kinesis.retrieval.GetRecordsCache;
@RunWith(MockitoJUnitRunner.class) @RunWith(MockitoJUnitRunner.class)
@ -72,7 +71,7 @@ public class ConsumerStatesTest {
@Mock @Mock
private ShardInfo shardInfo; private ShardInfo shardInfo;
@Mock @Mock
private LeaseManager<KinesisClientLease> leaseManager; private LeaseRefresher leaseRefresher;
@Mock @Mock
private Checkpointer checkpoint; private Checkpointer checkpoint;
@Mock @Mock
@ -84,7 +83,7 @@ public class ConsumerStatesTest {
@Mock @Mock
private AmazonKinesis amazonKinesis; private AmazonKinesis amazonKinesis;
@Mock @Mock
private LeaseManagerProxy leaseManagerProxy; private ShardDetector shardDetector;
@Mock @Mock
private IMetricsFactory metricsFactory; private IMetricsFactory metricsFactory;
@ -102,18 +101,18 @@ public class ConsumerStatesTest {
@Before @Before
public void setup() { public void setup() {
consumer = spy(new ShardConsumer(shardInfo, STREAM_NAME, leaseManager, executorService, getRecordsCache, consumer = spy(new ShardConsumer(shardInfo, STREAM_NAME, leaseRefresher, executorService, getRecordsCache,
recordProcessor, checkpoint, recordProcessorCheckpointer, parentShardPollIntervalMillis, recordProcessor, checkpoint, recordProcessorCheckpointer, parentShardPollIntervalMillis,
taskBackoffTimeMillis, Optional.empty(), amazonKinesis, taskBackoffTimeMillis, Optional.empty(), amazonKinesis,
skipShardSyncAtWorkerInitializationIfLeasesExist, listShardsBackoffTimeInMillis, skipShardSyncAtWorkerInitializationIfLeasesExist, listShardsBackoffTimeInMillis,
maxListShardsRetryAttempts, shouldCallProcessRecordsEvenForEmptyRecordList, idleTimeInMillis, maxListShardsRetryAttempts, shouldCallProcessRecordsEvenForEmptyRecordList, idleTimeInMillis,
INITIAL_POSITION_IN_STREAM, cleanupLeasesOfCompletedShards, ignoreUnexpectedChildShards, INITIAL_POSITION_IN_STREAM, cleanupLeasesOfCompletedShards, ignoreUnexpectedChildShards,
leaseManagerProxy, metricsFactory)); shardDetector, metricsFactory));
when(shardInfo.shardId()).thenReturn("shardId-000000000000"); when(shardInfo.shardId()).thenReturn("shardId-000000000000");
} }
private static final Class<LeaseManager<KinesisClientLease>> LEASE_MANAGER_CLASS = (Class<LeaseManager<KinesisClientLease>>) (Class<?>) LeaseManager.class; private static final Class<LeaseRefresher> LEASE_REFRESHER_CLASS = (Class<LeaseRefresher>) (Class<?>) LeaseRefresher.class;
@Test @Test
public void blockOnParentStateTest() { public void blockOnParentStateTest() {
@ -123,7 +122,7 @@ public class ConsumerStatesTest {
assertThat(task, taskWith(BlockOnParentShardTask.class, ShardInfo.class, "shardInfo", equalTo(shardInfo))); assertThat(task, taskWith(BlockOnParentShardTask.class, ShardInfo.class, "shardInfo", equalTo(shardInfo)));
assertThat(task, assertThat(task,
taskWith(BlockOnParentShardTask.class, LEASE_MANAGER_CLASS, "leaseManager", equalTo(leaseManager))); taskWith(BlockOnParentShardTask.class, LEASE_REFRESHER_CLASS, "leaseRefresher", equalTo(leaseRefresher)));
assertThat(task, taskWith(BlockOnParentShardTask.class, Long.class, "parentShardPollIntervalMillis", assertThat(task, taskWith(BlockOnParentShardTask.class, Long.class, "parentShardPollIntervalMillis",
equalTo(parentShardPollIntervalMillis))); equalTo(parentShardPollIntervalMillis)));
@ -300,7 +299,7 @@ public class ConsumerStatesTest {
assertThat(task, shutdownTask(RecordProcessorCheckpointer.class, "recordProcessorCheckpointer", assertThat(task, shutdownTask(RecordProcessorCheckpointer.class, "recordProcessorCheckpointer",
equalTo(recordProcessorCheckpointer))); equalTo(recordProcessorCheckpointer)));
assertThat(task, shutdownTask(ShutdownReason.class, "reason", equalTo(reason))); assertThat(task, shutdownTask(ShutdownReason.class, "reason", equalTo(reason)));
assertThat(task, shutdownTask(LEASE_MANAGER_CLASS, "leaseManager", equalTo(leaseManager))); assertThat(task, shutdownTask(LEASE_REFRESHER_CLASS, "leaseRefresher", equalTo(leaseRefresher)));
assertThat(task, shutdownTask(InitialPositionInStreamExtended.class, "initialPositionInStream", assertThat(task, shutdownTask(InitialPositionInStreamExtended.class, "initialPositionInStream",
equalTo(initialPositionInStream))); equalTo(initialPositionInStream)));
assertThat(task, assertThat(task,

View file

@ -47,7 +47,7 @@ import com.google.protobuf.ByteString;
import lombok.Data; import lombok.Data;
import software.amazon.kinesis.checkpoint.RecordProcessorCheckpointer; import software.amazon.kinesis.checkpoint.RecordProcessorCheckpointer;
import software.amazon.kinesis.leases.LeaseManagerProxy; import software.amazon.kinesis.leases.ShardDetector;
import software.amazon.kinesis.leases.ShardInfo; import software.amazon.kinesis.leases.ShardInfo;
import software.amazon.kinesis.processor.RecordProcessor; import software.amazon.kinesis.processor.RecordProcessor;
import software.amazon.kinesis.retrieval.ThrottlingReporter; import software.amazon.kinesis.retrieval.ThrottlingReporter;
@ -67,7 +67,7 @@ public class ProcessTaskTest {
@Mock @Mock
private ProcessRecordsInput processRecordsInput; private ProcessRecordsInput processRecordsInput;
@Mock @Mock
private LeaseManagerProxy leaseManagerProxy; private ShardDetector shardDetector;
@SuppressWarnings("serial") @SuppressWarnings("serial")
private static class RecordSubclass extends Record { private static class RecordSubclass extends Record {
@ -95,7 +95,7 @@ public class ProcessTaskTest {
private ProcessTask makeProcessTask(ProcessRecordsInput processRecordsInput) { private ProcessTask makeProcessTask(ProcessRecordsInput processRecordsInput) {
return new ProcessTask(shardInfo, recordProcessor, checkpointer, taskBackoffTimeMillis, return new ProcessTask(shardInfo, recordProcessor, checkpointer, taskBackoffTimeMillis,
skipShardSyncAtWorkerInitializationIfLeasesExist, leaseManagerProxy, throttlingReporter, skipShardSyncAtWorkerInitializationIfLeasesExist, shardDetector, throttlingReporter,
processRecordsInput, shouldCallProcessRecordsEvenForEmptyRecordList, IDLE_TIME_IN_MILLISECONDS); processRecordsInput, shouldCallProcessRecordsEvenForEmptyRecordList, IDLE_TIME_IN_MILLISECONDS);
} }

View file

@ -73,11 +73,10 @@ import com.amazonaws.services.kinesis.model.ShardIteratorType;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.checkpoint.Checkpoint; import software.amazon.kinesis.checkpoint.Checkpoint;
import software.amazon.kinesis.coordinator.KinesisClientLibConfiguration;
import software.amazon.kinesis.checkpoint.RecordProcessorCheckpointer; import software.amazon.kinesis.checkpoint.RecordProcessorCheckpointer;
import software.amazon.kinesis.leases.LeaseManager; import software.amazon.kinesis.coordinator.KinesisClientLibConfiguration;
import software.amazon.kinesis.leases.KinesisClientLease; import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.LeaseManagerProxy; import software.amazon.kinesis.leases.ShardDetector;
import software.amazon.kinesis.leases.ShardInfo; import software.amazon.kinesis.leases.ShardInfo;
import software.amazon.kinesis.metrics.IMetricsFactory; import software.amazon.kinesis.metrics.IMetricsFactory;
import software.amazon.kinesis.metrics.NullMetricsFactory; import software.amazon.kinesis.metrics.NullMetricsFactory;
@ -139,7 +138,7 @@ public class ShardConsumerTest {
@Mock @Mock
private KinesisClientLibConfiguration config; private KinesisClientLibConfiguration config;
@Mock @Mock
private LeaseManager<KinesisClientLease> leaseManager; private LeaseRefresher leaseRefresher;
@Mock @Mock
private Checkpointer checkpoint; private Checkpointer checkpoint;
@Mock @Mock
@ -149,7 +148,7 @@ public class ShardConsumerTest {
@Mock @Mock
private RecordProcessorCheckpointer recordProcessorCheckpointer; private RecordProcessorCheckpointer recordProcessorCheckpointer;
@Mock @Mock
private LeaseManagerProxy leaseManagerProxy; private ShardDetector shardDetector;
@Before @Before
public void setup() { public void setup() {
@ -225,7 +224,7 @@ public class ShardConsumerTest {
final ExtendedSequenceNumber checkpointSequenceNumber = new ExtendedSequenceNumber("123"); final ExtendedSequenceNumber checkpointSequenceNumber = new ExtendedSequenceNumber("123");
final ExtendedSequenceNumber pendingCheckpointSequenceNumber = null; final ExtendedSequenceNumber pendingCheckpointSequenceNumber = null;
when(leaseManager.getLease(anyString())).thenReturn(null); when(leaseRefresher.getLease(anyString())).thenReturn(null);
when(checkpoint.getCheckpointObject(anyString())).thenReturn( when(checkpoint.getCheckpointObject(anyString())).thenReturn(
new Checkpoint(checkpointSequenceNumber, pendingCheckpointSequenceNumber)); new Checkpoint(checkpointSequenceNumber, pendingCheckpointSequenceNumber));
@ -289,7 +288,7 @@ public class ShardConsumerTest {
final int maxRecords = 2; final int maxRecords = 2;
Checkpointer checkpoint = new InMemoryCheckpointer(startSeqNum.toString()); Checkpointer checkpoint = new InMemoryCheckpointer(startSeqNum.toString());
checkpoint.setCheckpoint(shardId, ExtendedSequenceNumber.TRIM_HORIZON, concurrencyToken); checkpoint.setCheckpoint(shardId, ExtendedSequenceNumber.TRIM_HORIZON, concurrencyToken);
when(leaseManager.getLease(anyString())).thenReturn(null); when(leaseRefresher.getLease(anyString())).thenReturn(null);
TestStreamlet processor = new TestStreamlet(); TestStreamlet processor = new TestStreamlet();
shardInfo = new ShardInfo(shardId, concurrencyToken, null, null); shardInfo = new ShardInfo(shardId, concurrencyToken, null, null);
@ -392,7 +391,7 @@ public class ShardConsumerTest {
final int maxRecords = 2; final int maxRecords = 2;
Checkpointer checkpoint = new InMemoryCheckpointer(startSeqNum.toString()); Checkpointer checkpoint = new InMemoryCheckpointer(startSeqNum.toString());
checkpoint.setCheckpoint(shardId, ExtendedSequenceNumber.TRIM_HORIZON, concurrencyToken); checkpoint.setCheckpoint(shardId, ExtendedSequenceNumber.TRIM_HORIZON, concurrencyToken);
when(leaseManager.getLease(anyString())).thenReturn(null); when(leaseRefresher.getLease(anyString())).thenReturn(null);
TransientShutdownErrorTestStreamlet processor = new TransientShutdownErrorTestStreamlet(); TransientShutdownErrorTestStreamlet processor = new TransientShutdownErrorTestStreamlet();
shardInfo = new ShardInfo(shardId, concurrencyToken, null, null); shardInfo = new ShardInfo(shardId, concurrencyToken, null, null);
@ -487,7 +486,7 @@ public class ShardConsumerTest {
final int maxRecords = 2; final int maxRecords = 2;
Checkpointer checkpoint = new InMemoryCheckpointer(startSeqNum.toString()); Checkpointer checkpoint = new InMemoryCheckpointer(startSeqNum.toString());
checkpoint.setCheckpoint(streamShardId, ExtendedSequenceNumber.AT_TIMESTAMP, testConcurrencyToken); checkpoint.setCheckpoint(streamShardId, ExtendedSequenceNumber.AT_TIMESTAMP, testConcurrencyToken);
when(leaseManager.getLease(anyString())).thenReturn(null); when(leaseRefresher.getLease(anyString())).thenReturn(null);
TestStreamlet processor = new TestStreamlet(); TestStreamlet processor = new TestStreamlet();
when(recordsFetcherFactory.createRecordsFetcher(any(GetRecordsRetrievalStrategy.class), anyString(), when(recordsFetcherFactory.createRecordsFetcher(any(GetRecordsRetrievalStrategy.class), anyString(),
@ -548,7 +547,7 @@ public class ShardConsumerTest {
final ExtendedSequenceNumber checkpointSequenceNumber = new ExtendedSequenceNumber("123"); final ExtendedSequenceNumber checkpointSequenceNumber = new ExtendedSequenceNumber("123");
final ExtendedSequenceNumber pendingCheckpointSequenceNumber = new ExtendedSequenceNumber("999"); final ExtendedSequenceNumber pendingCheckpointSequenceNumber = new ExtendedSequenceNumber("999");
when(leaseManager.getLease(anyString())).thenReturn(null); when(leaseRefresher.getLease(anyString())).thenReturn(null);
when(config.getRecordsFetcherFactory()).thenReturn(new SimpleRecordsFetcherFactory()); when(config.getRecordsFetcherFactory()).thenReturn(new SimpleRecordsFetcherFactory());
when(checkpoint.getCheckpointObject(anyString())).thenReturn( when(checkpoint.getCheckpointObject(anyString())).thenReturn(
new Checkpoint(checkpointSequenceNumber, pendingCheckpointSequenceNumber)); new Checkpoint(checkpointSequenceNumber, pendingCheckpointSequenceNumber));
@ -642,7 +641,7 @@ public class ShardConsumerTest {
final Optional<Long> logWarningForTaskAfterMillis) { final Optional<Long> logWarningForTaskAfterMillis) {
return new ShardConsumer(shardInfo, return new ShardConsumer(shardInfo,
streamName, streamName,
leaseManager, leaseRefresher,
executorService, executorService,
getRecordsCache, getRecordsCache,
recordProcessor, recordProcessor,
@ -660,7 +659,7 @@ public class ShardConsumerTest {
initialPositionLatest, initialPositionLatest,
cleanupLeasesOfCompletedShards, cleanupLeasesOfCompletedShards,
ignoreUnexpectedChildShards, ignoreUnexpectedChildShards,
leaseManagerProxy, shardDetector,
metricsFactory); metricsFactory);
} }

View file

@ -34,9 +34,8 @@ import com.amazonaws.services.kinesis.clientlibrary.lib.worker.InitialPositionIn
import com.amazonaws.services.kinesis.clientlibrary.lib.worker.InitialPositionInStreamExtended; import com.amazonaws.services.kinesis.clientlibrary.lib.worker.InitialPositionInStreamExtended;
import software.amazon.kinesis.checkpoint.RecordProcessorCheckpointer; import software.amazon.kinesis.checkpoint.RecordProcessorCheckpointer;
import software.amazon.kinesis.leases.LeaseManager; import software.amazon.kinesis.leases.LeaseRefresher;
import software.amazon.kinesis.leases.KinesisClientLease; import software.amazon.kinesis.leases.ShardDetector;
import software.amazon.kinesis.leases.LeaseManagerProxy;
import software.amazon.kinesis.leases.ShardInfo; import software.amazon.kinesis.leases.ShardInfo;
import software.amazon.kinesis.processor.RecordProcessor; import software.amazon.kinesis.processor.RecordProcessor;
import software.amazon.kinesis.retrieval.GetRecordsCache; import software.amazon.kinesis.retrieval.GetRecordsCache;
@ -66,9 +65,9 @@ public class ShutdownTaskTest {
@Mock @Mock
private RecordProcessorCheckpointer checkpointer; private RecordProcessorCheckpointer checkpointer;
@Mock @Mock
private LeaseManager<KinesisClientLease> leaseManager; private LeaseRefresher leaseRefresher;
@Mock @Mock
private LeaseManagerProxy leaseManagerProxy; private ShardDetector shardDetector;
@Before @Before
public void setUp() throws Exception { public void setUp() throws Exception {
@ -78,9 +77,9 @@ public class ShutdownTaskTest {
ExtendedSequenceNumber.LATEST); ExtendedSequenceNumber.LATEST);
recordProcessor = new TestStreamlet(); recordProcessor = new TestStreamlet();
task = new ShutdownTask(shardInfo, leaseManagerProxy, recordProcessor, checkpointer, task = new ShutdownTask(shardInfo, shardDetector, recordProcessor, checkpointer,
TERMINATE_SHUTDOWN_REASON, INITIAL_POSITION_TRIM_HORIZON, cleanupLeasesOfCompletedShards, TERMINATE_SHUTDOWN_REASON, INITIAL_POSITION_TRIM_HORIZON, cleanupLeasesOfCompletedShards,
ignoreUnexpectedChildShards, leaseManager, TASK_BACKOFF_TIME_MILLIS, getRecordsCache); ignoreUnexpectedChildShards, leaseRefresher, TASK_BACKOFF_TIME_MILLIS, getRecordsCache);
} }
/** /**
@ -100,7 +99,7 @@ public class ShutdownTaskTest {
@Test @Test
public final void testCallWhenSyncingShardsThrows() { public final void testCallWhenSyncingShardsThrows() {
when(checkpointer.lastCheckpointValue()).thenReturn(ExtendedSequenceNumber.SHARD_END); when(checkpointer.lastCheckpointValue()).thenReturn(ExtendedSequenceNumber.SHARD_END);
when(leaseManagerProxy.listShards()).thenReturn(null); when(shardDetector.listShards()).thenReturn(null);
TaskResult result = task.call(); TaskResult result = task.call();
assertNotNull(result.getException()); assertNotNull(result.getException());