Add config to enable PITR (#1365)

* Add config to enable PITR
This commit is contained in:
lucienlu-aws 2024-07-15 14:15:52 -07:00 committed by GitHub
parent 5878ba8ac6
commit 1c0c41c4e8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 268 additions and 12 deletions

View file

@ -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;

View file

@ -88,6 +88,12 @@ public class MultiLangDaemonConfiguration {
@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;

View file

@ -33,6 +33,7 @@ import software.amazon.kinesis.retrieval.polling.PollingConfig;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.instanceOf;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;
@ -99,6 +100,50 @@ public class MultiLangDaemonConfigurationTest {
assertThat(resolvedConfiguration.leaseManagementConfig.enablePriorityLeaseAssignment(), equalTo(false));
}
@Test
public void testSetLeaseTableDeletionProtectionEnabledToTrue() {
MultiLangDaemonConfiguration configuration = baseConfiguration();
configuration.setLeaseTableDeletionProtectionEnabled(true);
MultiLangDaemonConfiguration.ResolvedConfiguration resolvedConfiguration =
configuration.resolvedConfiguration(shardRecordProcessorFactory);
assertTrue(resolvedConfiguration.leaseManagementConfig.leaseTableDeletionProtectionEnabled());
}
@Test
public void testSetLeaseTablePitrEnabledToTrue() {
MultiLangDaemonConfiguration configuration = baseConfiguration();
configuration.setLeaseTablePitrEnabled(true);
MultiLangDaemonConfiguration.ResolvedConfiguration resolvedConfiguration =
configuration.resolvedConfiguration(shardRecordProcessorFactory);
assertTrue(resolvedConfiguration.leaseManagementConfig.leaseTablePitrEnabled());
}
@Test
public void testSetLeaseTableDeletionProtectionEnabledToFalse() {
MultiLangDaemonConfiguration configuration = baseConfiguration();
configuration.setLeaseTableDeletionProtectionEnabled(false);
MultiLangDaemonConfiguration.ResolvedConfiguration resolvedConfiguration =
configuration.resolvedConfiguration(shardRecordProcessorFactory);
assertFalse(resolvedConfiguration.leaseManagementConfig.leaseTableDeletionProtectionEnabled());
}
@Test
public void testSetLeaseTablePitrEnabledToFalse() {
MultiLangDaemonConfiguration configuration = baseConfiguration();
configuration.setLeaseTablePitrEnabled(false);
MultiLangDaemonConfiguration.ResolvedConfiguration resolvedConfiguration =
configuration.resolvedConfiguration(shardRecordProcessorFactory);
assertFalse(resolvedConfiguration.leaseManagementConfig.leaseTablePitrEnabled());
}
@Test
public void testDefaultRetrievalConfig() {
MultiLangDaemonConfiguration configuration = baseConfiguration();

View file

@ -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;
@ -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.
@ -424,6 +435,7 @@ public class LeaseManagementConfig {
dynamoDbRequestTimeout(),
billingMode(),
leaseTableDeletionProtectionEnabled(),
leaseTablePitrEnabled(),
tags(),
leaseSerializer,
customShardDetectorProvider(),

View file

@ -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;
@ -707,7 +708,7 @@ public class DynamoDBLeaseManagementFactory implements LeaseManagementFactory {
tableCreatorCallback,
dynamoDbRequestTimeout,
billingMode,
false,
LeaseManagementConfig.DEFAULT_LEASE_TABLE_DELETION_PROTECTION_ENABLED,
DefaultSdkAutoConstructList.getInstance(),
leaseSerializer);
}
@ -945,6 +946,7 @@ public class DynamoDBLeaseManagementFactory implements LeaseManagementFactory {
* @param isMultiStreamMode
* @param leaseCleanupConfig
*/
@Deprecated
public DynamoDBLeaseManagementFactory(
final KinesisAsyncClient kinesisClient,
final DynamoDbAsyncClient dynamoDBClient,
@ -978,6 +980,76 @@ public class DynamoDBLeaseManagementFactory implements LeaseManagementFactory {
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,
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,
final boolean leaseTablePitrEnabled,
Collection<Tag> tags,
LeaseSerializer leaseSerializer,
Function<StreamConfig, ShardDetector> customShardDetectorProvider,
boolean isMultiStreamMode,
LeaseCleanupConfig leaseCleanupConfig) {
this.kinesisClient = kinesisClient;
this.dynamoDBClient = dynamoDBClient;
this.tableName = tableName;
@ -1005,6 +1077,7 @@ public class DynamoDBLeaseManagementFactory implements LeaseManagementFactory {
this.dynamoDbRequestTimeout = dynamoDbRequestTimeout;
this.billingMode = billingMode;
this.leaseTableDeletionProtectionEnabled = leaseTableDeletionProtectionEnabled;
this.leaseTablePitrEnabled = leaseTablePitrEnabled;
this.leaseSerializer = leaseSerializer;
this.customShardDetectorProvider = customShardDetectorProvider;
this.isMultiStreamMode = isMultiStreamMode;
@ -1091,6 +1164,7 @@ public class DynamoDBLeaseManagementFactory implements LeaseManagementFactory {
dynamoDbRequestTimeout,
billingMode,
leaseTableDeletionProtectionEnabled,
leaseTablePitrEnabled,
tags);
}

View file

@ -49,6 +49,7 @@ import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
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;
@ -81,6 +82,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;
@ -159,7 +161,7 @@ public class DynamoDBLeaseRefresher implements LeaseRefresher {
tableCreatorCallback,
dynamoDbRequestTimeout,
BillingMode.PAY_PER_REQUEST,
false);
LeaseManagementConfig.DEFAULT_LEASE_TABLE_DELETION_PROTECTION_ENABLED);
}
/**
@ -207,6 +209,7 @@ public class DynamoDBLeaseRefresher implements LeaseRefresher {
* @param leaseTableDeletionProtectionEnabled
* @param tags
*/
@Deprecated
public DynamoDBLeaseRefresher(
final String table,
final DynamoDbAsyncClient dynamoDBClient,
@ -217,6 +220,41 @@ public class DynamoDBLeaseRefresher implements LeaseRefresher {
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,
final LeaseSerializer serializer,
final boolean consistentReads,
@NonNull final TableCreatorCallback tableCreatorCallback,
Duration dynamoDbRequestTimeout,
final BillingMode billingMode,
final boolean leaseTableDeletionProtectionEnabled,
final boolean leaseTablePitrEnabled,
final Collection<Tag> tags) {
this.table = table;
this.dynamoDBClient = dynamoDBClient;
this.serializer = serializer;
@ -225,6 +263,7 @@ public class DynamoDBLeaseRefresher implements LeaseRefresher {
this.dynamoDbRequestTimeout = dynamoDbRequestTimeout;
this.billingMode = billingMode;
this.leaseTableDeletionProtectionEnabled = leaseTableDeletionProtectionEnabled;
this.leaseTablePitrEnabled = leaseTablePitrEnabled;
this.tags = tags;
}
@ -252,7 +291,33 @@ public class DynamoDBLeaseRefresher implements LeaseRefresher {
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(builder -> builder.pointInTimeRecoveryEnabled(true))
.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)

View file

@ -31,6 +31,7 @@ import org.junit.rules.ExpectedException;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.runners.MockitoJUnitRunner;
import software.amazon.awssdk.core.util.DefaultSdkAutoConstructList;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.BillingMode;
@ -54,6 +55,8 @@ import software.amazon.awssdk.services.dynamodb.model.ScanResponse;
import software.amazon.awssdk.services.dynamodb.model.TableDescription;
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.UpdateContinuousBackupsResponse;
import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
import software.amazon.awssdk.services.dynamodb.model.UpdateItemResponse;
import software.amazon.kinesis.leases.Lease;
@ -80,6 +83,10 @@ public class DynamoDBLeaseRefresherTest {
private static final String TABLE_NAME = "test";
private static final boolean CONSISTENT_READS = true;
private static final boolean DELETION_PROTECTION_ENABLED = false;
private static final boolean PITR_ENABLED = true;
private static final Collection<Tag> EMPTY_TAGS = DefaultSdkAutoConstructList.getInstance();
private static final Collection<Tag> TAGS =
Collections.singletonList(Tag.builder().key("foo").value("bar").build());
@Mock
private DynamoDbAsyncClient dynamoDbClient;
@ -111,6 +118,9 @@ public class DynamoDBLeaseRefresherTest {
@Mock
private CompletableFuture<CreateTableResponse> mockCreateTableFuture;
@Mock
private CompletableFuture<UpdateContinuousBackupsResponse> mockUpdateContinuousBackupsFuture;
@Mock
private Lease lease;
@ -120,8 +130,7 @@ public class DynamoDBLeaseRefresherTest {
private DynamoDBLeaseRefresher leaseRefresher;
private DescribeTableRequest describeTableRequest;
private CreateTableRequest createTableRequest;
private Collection<Tag> tags;
private UpdateContinuousBackupsRequest updateContinuousBackupsRequest;
private Map<String, AttributeValue> serializedLease;
@Before
@ -139,6 +148,10 @@ public class DynamoDBLeaseRefresherTest {
.billingMode(BillingMode.PAY_PER_REQUEST)
.deletionProtectionEnabled(DELETION_PROTECTION_ENABLED)
.build();
updateContinuousBackupsRequest = UpdateContinuousBackupsRequest.builder()
.tableName(TABLE_NAME)
.pointInTimeRecoverySpecification(builder -> builder.pointInTimeRecoveryEnabled(PITR_ENABLED))
.build();
}
@Test
@ -353,7 +366,6 @@ public class DynamoDBLeaseRefresherTest {
@Test
public void testCreateLeaseTableWithTagsIfNotExists() throws Exception {
tags = Collections.singletonList(Tag.builder().key("foo").value("bar").build());
leaseRefresher = new DynamoDBLeaseRefresher(
TABLE_NAME,
dynamoDbClient,
@ -363,7 +375,7 @@ public class DynamoDBLeaseRefresherTest {
LeaseManagementConfig.DEFAULT_REQUEST_TIMEOUT,
BillingMode.PROVISIONED,
DELETION_PROTECTION_ENABLED,
tags);
TAGS);
when(dynamoDbClient.describeTable(describeTableRequest)).thenReturn(mockDescribeTableFuture);
when(mockDescribeTableFuture.get(
@ -382,7 +394,7 @@ public class DynamoDBLeaseRefresherTest {
.attributeDefinitions(leaseSerializer.getAttributeDefinitions())
.provisionedThroughput(throughput)
.deletionProtectionEnabled(DELETION_PROTECTION_ENABLED)
.tags(tags)
.tags(TAGS)
.build();
when(dynamoDbClient.createTable(createTableRequest)).thenReturn(mockCreateTableFuture);
when(mockCreateTableFuture.get(LeaseManagementConfig.DEFAULT_REQUEST_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS))
@ -423,9 +435,49 @@ public class DynamoDBLeaseRefresherTest {
Assert.assertTrue(result);
}
@Test
public void testCreateLeaseTableIfNotExistsWithPitrEnabled() throws Exception {
DynamoDBLeaseRefresher leaseRefresherWithEnabledPitr = new DynamoDBLeaseRefresher(
TABLE_NAME,
dynamoDbClient,
leaseSerializer,
CONSISTENT_READS,
tableCreatorCallback,
LeaseManagementConfig.DEFAULT_REQUEST_TIMEOUT,
BillingMode.PAY_PER_REQUEST,
DELETION_PROTECTION_ENABLED,
PITR_ENABLED,
EMPTY_TAGS);
when(dynamoDbClient.describeTable(describeTableRequest)).thenReturn(mockDescribeTableFuture);
when(mockDescribeTableFuture.get(
eq(LeaseManagementConfig.DEFAULT_REQUEST_TIMEOUT.toMillis()), eq(TimeUnit.MILLISECONDS)))
.thenThrow(ResourceNotFoundException.builder()
.message("Table doesn't exist")
.build());
when(dynamoDbClient.createTable(createTableRequest)).thenReturn(mockCreateTableFuture);
when(mockCreateTableFuture.get(
eq(LeaseManagementConfig.DEFAULT_REQUEST_TIMEOUT.toMillis()), eq(TimeUnit.MILLISECONDS)))
.thenReturn(null);
when(dynamoDbClient.updateContinuousBackups(updateContinuousBackupsRequest))
.thenReturn(mockUpdateContinuousBackupsFuture);
when(mockUpdateContinuousBackupsFuture.get(
eq(LeaseManagementConfig.DEFAULT_REQUEST_TIMEOUT.toMillis()), eq(TimeUnit.MILLISECONDS)))
.thenReturn(null);
final boolean result = leaseRefresherWithEnabledPitr.createLeaseTableIfNotExists();
verify(dynamoDbClient, times(1)).describeTable(describeTableRequest);
verify(dynamoDbClient, times(1)).createTable(createTableRequest);
verify(dynamoDbClient, times(1)).updateContinuousBackups(updateContinuousBackupsRequest);
verify(mockDescribeTableFuture, times(1))
.get(eq(LeaseManagementConfig.DEFAULT_REQUEST_TIMEOUT.toMillis()), eq(TimeUnit.MILLISECONDS));
verify(mockCreateTableFuture, times(1))
.get(eq(LeaseManagementConfig.DEFAULT_REQUEST_TIMEOUT.toMillis()), eq(TimeUnit.MILLISECONDS));
Assert.assertTrue(result);
}
@Test
public void testCreateLeaseTableProvisionedWithDeletionProtectionIfNotExists() throws Exception {
leaseRefresher = new DynamoDBLeaseRefresher(
DynamoDBLeaseRefresher leaseRefresherWithEnabledDeletionProtection = new DynamoDBLeaseRefresher(
TABLE_NAME,
dynamoDbClient,
leaseSerializer,
@ -458,7 +510,7 @@ public class DynamoDBLeaseRefresherTest {
eq(LeaseManagementConfig.DEFAULT_REQUEST_TIMEOUT.toMillis()), eq(TimeUnit.MILLISECONDS)))
.thenReturn(null);
final boolean result = leaseRefresher.createLeaseTableIfNotExists(10L, 10L);
final boolean result = leaseRefresherWithEnabledDeletionProtection.createLeaseTableIfNotExists(10L, 10L);
verify(dynamoDbClient, times(1)).describeTable(describeTableRequest);
verify(dynamoDbClient, times(1)).createTable(createTableRequest);