Skip to content

Commit

Permalink
Add config to enable PITR
Browse files Browse the repository at this point in the history
  • Loading branch information
@lucienlu-aws
lucienlu-aws committed Jul 11, 2024
1 parent f507066 commit 103961c
Show file tree
Hide file tree
Showing 7 changed files with 251 additions and 12 deletions.
2 changes: 2 additions & 0 deletions ...lang/src/main/java/software/amazon/kinesis/coordinator/KinesisClientLibConfiguration.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -218,6 +218,8 @@ public class KinesisClientLibConfiguration {
private AwsCredentialsProvider cloudWatchCredentialsProvider;
private long failoverTimeMillis;
private boolean enablePriorityLeaseAssignment;
private boolean leaseTableDeletionProtectionEnabled;
private boolean leaseTablePitrEnabled;
private String workerIdentifier;
private long shardSyncIntervalMillis;
private int maxRecords;
Expand Down
6 changes: 6 additions & 0 deletions .../src/main/java/software/amazon/kinesis/multilang/config/MultiLangDaemonConfiguration.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,12 @@ public void setWorkerId(String workerId) {
@ConfigurationSettable(configurationClass = LeaseManagementConfig.class)
private Boolean enablePriorityLeaseAssignment;

@ConfigurationSettable(configurationClass = LeaseManagementConfig.class)
private Boolean leaseTableDeletionProtectionEnabled;

@ConfigurationSettable(configurationClass = LeaseManagementConfig.class)
private Boolean leaseTablePitrEnabled;

@ConfigurationSettable(configurationClass = LeaseManagementConfig.class)
private long shardSyncIntervalMillis;

Expand Down
22 changes: 22 additions & 0 deletions .../test/java/software/amazon/kinesis/multilang/config/MultiLangDaemonConfigurationTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,28 @@ public void testSetEnablePriorityLeaseAssignment() {
assertThat(resolvedConfiguration.leaseManagementConfig.enablePriorityLeaseAssignment(), equalTo(false));
}

@Test
public void testSetLeaseTableDeletionProtectionEnabled() {
MultiLangDaemonConfiguration configuration = baseConfiguration();
configuration.setLeaseTableDeletionProtectionEnabled(true);

MultiLangDaemonConfiguration.ResolvedConfiguration resolvedConfiguration =
configuration.resolvedConfiguration(shardRecordProcessorFactory);

assertThat(resolvedConfiguration.leaseManagementConfig.leaseTableDeletionProtectionEnabled(), equalTo(true));
}

@Test
public void testSetLeaseTablePitrEnabled() {
MultiLangDaemonConfiguration configuration = baseConfiguration();
configuration.setLeaseTablePitrEnabled(true);

MultiLangDaemonConfiguration.ResolvedConfiguration resolvedConfiguration =
configuration.resolvedConfiguration(shardRecordProcessorFactory);

assertThat(resolvedConfiguration.leaseManagementConfig.leaseTablePitrEnabled(), equalTo(true));
}

@Test
public void testDefaultRetrievalConfig() {
MultiLangDaemonConfiguration configuration = baseConfiguration();
Expand Down
16 changes: 14 additions & 2 deletions ...on-kinesis-client/src/main/java/software/amazon/kinesis/leases/LeaseManagementConfig.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,8 @@ public class LeaseManagementConfig {
public static final long DEFAULT_GARBAGE_LEASE_CLEANUP_INTERVAL_MILLIS =
Duration.ofMinutes(30).toMillis();
public static final long DEFAULT_PERIODIC_SHARD_SYNC_INTERVAL_MILLIS = 2 * 60 * 1000L;
public static final boolean DEFAULT_LEASE_TABLE_DELETION_PROTECTION_ENABLED = false;
public static final boolean DEFAULT_LEASE_TABLE_PITR_ENABLED = false;
public static final boolean DEFAULT_ENABLE_PRIORITY_LEASE_ASSIGNMENT = true;
public static final int DEFAULT_CONSECUTIVE_HOLES_FOR_TRIGGERING_LEASE_RECOVERY = 3;

Expand Down Expand Up @@ -208,11 +210,20 @@ public class LeaseManagementConfig {
private BillingMode billingMode = BillingMode.PAY_PER_REQUEST;

/**
* Whether to enabled deletion protection on the DynamoDB lease table created by KCL.
* Whether to enable deletion protection on the DynamoDB lease table created by KCL. This does not update
* already existing tables.
*
* <p>Default value: false
*/
private boolean leaseTableDeletionProtectionEnabled = false;
private boolean leaseTableDeletionProtectionEnabled = DEFAULT_LEASE_TABLE_DELETION_PROTECTION_ENABLED;

/**
* Whether to enable PITR (point in time recovery) on the DynamoDB lease table created by KCL. If true, this can
* update existing table's PITR.
*
* <p>Default value: false
*/
private boolean leaseTablePitrEnabled = DEFAULT_LEASE_TABLE_PITR_ENABLED;

/**
* The list of tags to be applied to the DynamoDB table created for lease management.
Expand Down Expand Up @@ -424,6 +435,7 @@ public LeaseManagementFactory leaseManagementFactory(
dynamoDbRequestTimeout(),
billingMode(),
leaseTableDeletionProtectionEnabled(),
leaseTablePitrEnabled(),
tags(),
leaseSerializer,
customShardDetectorProvider(),
Expand Down
76 changes: 75 additions & 1 deletion ...src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseManagementFactory.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,7 @@ public class DynamoDBLeaseManagementFactory implements LeaseManagementFactory {
private final Duration dynamoDbRequestTimeout;
private final BillingMode billingMode;
private final boolean leaseTableDeletionProtectionEnabled;
private final boolean leaseTablePitrEnabled;
private final Collection<Tag> tags;
private final boolean isMultiStreamMode;
private final LeaseCleanupConfig leaseCleanupConfig;
Expand Down Expand Up @@ -707,7 +708,7 @@ private DynamoDBLeaseManagementFactory(
tableCreatorCallback,
dynamoDbRequestTimeout,
billingMode,
false,
LeaseManagementConfig.DEFAULT_LEASE_TABLE_DELETION_PROTECTION_ENABLED,
DefaultSdkAutoConstructList.getInstance(),
leaseSerializer);
}
Expand Down Expand Up @@ -945,6 +946,76 @@ public DynamoDBLeaseManagementFactory(
* @param isMultiStreamMode
* @param leaseCleanupConfig
*/
@Deprecated
public DynamoDBLeaseManagementFactory(
final KinesisAsyncClient kinesisClient,
final DynamoDbAsyncClient dynamoDBClient,
final String tableName,
final String workerIdentifier,
final ExecutorService executorService,
final long failoverTimeMillis,
final boolean enablePriorityLeaseAssignment,
final long epsilonMillis,
final int maxLeasesForWorker,
final int maxLeasesToStealAtOneTime,
final int maxLeaseRenewalThreads,
final boolean cleanupLeasesUponShardCompletion,
final boolean ignoreUnexpectedChildShards,
final long shardSyncIntervalMillis,
final boolean consistentReads,
final long listShardsBackoffTimeMillis,
final int maxListShardsRetryAttempts,
final int maxCacheMissesBeforeReload,
final long listShardsCacheAllowedAgeInSeconds,
final int cacheMissWarningModulus,
final long initialLeaseTableReadCapacity,
final long initialLeaseTableWriteCapacity,
final HierarchicalShardSyncer deprecatedHierarchicalShardSyncer,
final TableCreatorCallback tableCreatorCallback,
Duration dynamoDbRequestTimeout,
BillingMode billingMode,
final boolean leaseTableDeletionProtectionEnabled,
Collection<Tag> tags,
LeaseSerializer leaseSerializer,
Function<StreamConfig, ShardDetector> customShardDetectorProvider,
boolean isMultiStreamMode,
LeaseCleanupConfig leaseCleanupConfig) {
this(
kinesisClient,
dynamoDBClient,
tableName,
workerIdentifier,
executorService,
failoverTimeMillis,
enablePriorityLeaseAssignment,
epsilonMillis,
maxLeasesForWorker,
maxLeasesToStealAtOneTime,
maxLeaseRenewalThreads,
cleanupLeasesUponShardCompletion,
ignoreUnexpectedChildShards,
shardSyncIntervalMillis,
consistentReads,
listShardsBackoffTimeMillis,
maxListShardsRetryAttempts,
maxCacheMissesBeforeReload,
listShardsCacheAllowedAgeInSeconds,
cacheMissWarningModulus,
initialLeaseTableReadCapacity,
initialLeaseTableWriteCapacity,
deprecatedHierarchicalShardSyncer,
tableCreatorCallback,
dynamoDbRequestTimeout,
billingMode,
leaseTableDeletionProtectionEnabled,
LeaseManagementConfig.DEFAULT_LEASE_TABLE_PITR_ENABLED,
tags,
leaseSerializer,
customShardDetectorProvider,
isMultiStreamMode,
leaseCleanupConfig);
}

public DynamoDBLeaseManagementFactory(
final KinesisAsyncClient kinesisClient,
final DynamoDbAsyncClient dynamoDBClient,
Expand Down Expand Up @@ -973,6 +1044,7 @@ public DynamoDBLeaseManagementFactory(
Duration dynamoDbRequestTimeout,
BillingMode billingMode,
final boolean leaseTableDeletionProtectionEnabled,
final boolean leaseTablePitrEnabled,
Collection<Tag> tags,
LeaseSerializer leaseSerializer,
Function<StreamConfig, ShardDetector> customShardDetectorProvider,
Expand Down Expand Up @@ -1005,6 +1077,7 @@ public DynamoDBLeaseManagementFactory(
this.dynamoDbRequestTimeout = dynamoDbRequestTimeout;
this.billingMode = billingMode;
this.leaseTableDeletionProtectionEnabled = leaseTableDeletionProtectionEnabled;
this.leaseTablePitrEnabled = leaseTablePitrEnabled;
this.leaseSerializer = leaseSerializer;
this.customShardDetectorProvider = customShardDetectorProvider;
this.isMultiStreamMode = isMultiStreamMode;
Expand Down Expand Up @@ -1091,6 +1164,7 @@ public DynamoDBLeaseRefresher createLeaseRefresher() {
dynamoDbRequestTimeout,
billingMode,
leaseTableDeletionProtectionEnabled,
leaseTablePitrEnabled,
tags);
}

Expand Down
72 changes: 70 additions & 2 deletions ...-client/src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseRefresher.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
import software.amazon.awssdk.services.dynamodb.model.LimitExceededException;
import software.amazon.awssdk.services.dynamodb.model.PointInTimeRecoverySpecification;
import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughput;
import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughputExceededException;
import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
Expand All @@ -49,6 +50,7 @@
import software.amazon.awssdk.services.dynamodb.model.ScanResponse;
import software.amazon.awssdk.services.dynamodb.model.TableStatus;
import software.amazon.awssdk.services.dynamodb.model.Tag;
import software.amazon.awssdk.services.dynamodb.model.UpdateContinuousBackupsRequest;
import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
import software.amazon.awssdk.utils.CollectionUtils;
import software.amazon.kinesis.annotations.KinesisClientInternalApi;
Expand Down Expand Up @@ -81,6 +83,7 @@ public class DynamoDBLeaseRefresher implements LeaseRefresher {
private final Duration dynamoDbRequestTimeout;
private final BillingMode billingMode;
private final boolean leaseTableDeletionProtectionEnabled;
private final boolean leaseTablePitrEnabled;
private final Collection<Tag> tags;

private boolean newTableCreated = false;
Expand Down Expand Up @@ -159,7 +162,7 @@ public DynamoDBLeaseRefresher(
tableCreatorCallback,
dynamoDbRequestTimeout,
BillingMode.PAY_PER_REQUEST,
false);
LeaseManagementConfig.DEFAULT_LEASE_TABLE_DELETION_PROTECTION_ENABLED);
}

/**
Expand Down Expand Up @@ -207,6 +210,41 @@ public DynamoDBLeaseRefresher(
* @param leaseTableDeletionProtectionEnabled
* @param tags
*/
@Deprecated
public DynamoDBLeaseRefresher(
final String table,
final DynamoDbAsyncClient dynamoDBClient,
final LeaseSerializer serializer,
final boolean consistentReads,
@NonNull final TableCreatorCallback tableCreatorCallback,
Duration dynamoDbRequestTimeout,
final BillingMode billingMode,
final boolean leaseTableDeletionProtectionEnabled,
final Collection<Tag> tags) {
this(
table,
dynamoDBClient,
serializer,
consistentReads,
tableCreatorCallback,
dynamoDbRequestTimeout,
billingMode,
leaseTableDeletionProtectionEnabled,
LeaseManagementConfig.DEFAULT_LEASE_TABLE_PITR_ENABLED,
tags);
}

/**
* Constructor.
* @param table
* @param dynamoDBClient
* @param serializer
* @param consistentReads
* @param tableCreatorCallback
* @param dynamoDbRequestTimeout
* @param billingMode
* @param leaseTableDeletionProtectionEnabled
*/
public DynamoDBLeaseRefresher(
final String table,
final DynamoDbAsyncClient dynamoDBClient,
Expand All @@ -216,6 +254,7 @@ public DynamoDBLeaseRefresher(
Duration dynamoDbRequestTimeout,
final BillingMode billingMode,
final boolean leaseTableDeletionProtectionEnabled,
final boolean leaseTablePitrEnabled,
final Collection<Tag> tags) {
this.table = table;
this.dynamoDBClient = dynamoDBClient;
Expand All @@ -225,6 +264,7 @@ public DynamoDBLeaseRefresher(
this.dynamoDbRequestTimeout = dynamoDbRequestTimeout;
this.billingMode = billingMode;
this.leaseTableDeletionProtectionEnabled = leaseTableDeletionProtectionEnabled;
this.leaseTablePitrEnabled = leaseTablePitrEnabled;
this.tags = tags;
}

Expand Down Expand Up @@ -252,7 +292,35 @@ public boolean createLeaseTableIfNotExists(@NonNull final Long readCapacity, @No
public boolean createLeaseTableIfNotExists() throws ProvisionedThroughputException, DependencyException {
final CreateTableRequest request = createTableRequestBuilder().build();

return createTableIfNotExists(request);
boolean tableExists = createTableIfNotExists(request);

if (leaseTablePitrEnabled) {
enablePitr();
log.info("Enabled PITR on table {}", table);
}

return tableExists;
}

private void enablePitr() throws DependencyException {
final UpdateContinuousBackupsRequest request = UpdateContinuousBackupsRequest.builder()
.tableName(table)
.pointInTimeRecoverySpecification(PointInTimeRecoverySpecification.builder()
.pointInTimeRecoveryEnabled(true)
.build())
.build();

final AWSExceptionManager exceptionManager = createExceptionManager();
exceptionManager.add(ResourceNotFoundException.class, t -> t);
exceptionManager.add(ProvisionedThroughputExceededException.class, t -> t);

try {
FutureUtils.resolveOrCancelFuture(dynamoDBClient.updateContinuousBackups(request), dynamoDbRequestTimeout);
} catch (ExecutionException e) {
throw exceptionManager.apply(e.getCause());
} catch (InterruptedException | DynamoDbException | TimeoutException e) {
throw new DependencyException(e);
}
}

private boolean createTableIfNotExists(CreateTableRequest request)
Expand Down
Loading

0 comments on commit 103961c

Please sign in to comment.