Extend ttl when items are updated to keep items that are being used f…

…rom expiring Signed-off-by: Taylor Gray <[email protected]>
opensearch-project · Aug 8, 2023 · b18516d · b18516d
1 parent 4ed52b3
commit b18516d
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/...earch/dataprepper/plugins/sourcecoordinator/dynamodb/DynamoDbSourceCoordinationStore.java b/...earch/dataprepper/plugins/sourcecoordinator/dynamodb/DynamoDbSourceCoordinationStore.java
@@ -118,6 +118,10 @@ public void tryUpdateSourcePartitionItem(final SourcePartitionStoreItem updateIt
             dynamoDbSourcePartitionItem.setPartitionPriority(updateItem.getPartitionOwnershipTimeout().toString());
         }
 
+        if (Objects.nonNull(dynamoStoreSettings.getTtl())) {
+            dynamoDbSourcePartitionItem.setExpirationTime(Instant.now().plus(dynamoStoreSettings.getTtl()).getEpochSecond());
+        }
+
         dynamoDbClientWrapper.tryUpdatePartitionItem(dynamoDbSourcePartitionItem);
     }
 

diff --git a/...h/dataprepper/plugins/sourcecoordinator/dynamodb/DynamoDbSourceCoordinationStoreTest.java b/...h/dataprepper/plugins/sourcecoordinator/dynamodb/DynamoDbSourceCoordinationStoreTest.java
@@ -133,6 +133,8 @@ void tryUpdateSourcePartitionItem_calls_dynamoClientWrapper_correctly_for_assign
         given(updateItem.getSourceIdentifier()).willReturn(sourceIdentifier);
         given(updateItem.getSourcePartitionStatus()).willReturn(SourcePartitionStatus.ASSIGNED);
 
+        given(dynamoStoreSettings.getTtl()).willReturn(null);
+
         doNothing().when(dynamoDbClientWrapper).tryUpdatePartitionItem((DynamoDbSourcePartitionItem) updateItem);
 
         final Instant partitionOwnershipTimeout = Instant.now();
@@ -151,6 +153,8 @@ void tryUpdateSourcePartitionItem_calls_dynamoClientWrapper_correctly_for_closed
         given(updateItem.getSourceIdentifier()).willReturn(sourceIdentifier);
         given(updateItem.getSourcePartitionStatus()).willReturn(SourcePartitionStatus.CLOSED);
 
+        given(dynamoStoreSettings.getTtl()).willReturn(null);
+
         doNothing().when(dynamoDbClientWrapper).tryUpdatePartitionItem((DynamoDbSourcePartitionItem) updateItem);
         final Instant reOpenAtTime = Instant.now();
         given(updateItem.getReOpenAt()).willReturn(reOpenAtTime);
@@ -168,10 +172,19 @@ void tryUpdateSourcePartitionItem_calls_dynamoClientWrapper_correctly_for_comple
         given(updateItem.getSourceIdentifier()).willReturn(sourceIdentifier);
         given(updateItem.getSourcePartitionStatus()).willReturn(SourcePartitionStatus.COMPLETED);
 
+        final ArgumentCaptor<Long> argumentCaptor = ArgumentCaptor.forClass(Long.class);
+        final Duration ttl = Duration.ofSeconds(30);
+        final Long nowPlusTtl = Instant.now().plus(ttl).getEpochSecond();
+        given(dynamoStoreSettings.getTtl()).willReturn(ttl);
+        doNothing().when((DynamoDbSourcePartitionItem) updateItem).setExpirationTime(argumentCaptor.capture());
+
         doNothing().when(dynamoDbClientWrapper).tryUpdatePartitionItem((DynamoDbSourcePartitionItem) updateItem);
 
         createObjectUnderTest().tryUpdateSourcePartitionItem(updateItem);
 
+        final Long expirationTimeResult = argumentCaptor.getValue();
+        assertThat(expirationTimeResult, greaterThanOrEqualTo(nowPlusTtl));
+
         verify((DynamoDbSourcePartitionItem) updateItem).setSourceStatusCombinationKey(sourceIdentifier + "|" + SourcePartitionStatus.COMPLETED);
         verify((DynamoDbSourcePartitionItem) updateItem, never()).setPartitionPriority(anyString());
     }