Skip to content

Commit fbd6754

Browse files
author
ebartkus
committed
KCLRecordProcessor handles multiple shutdown calls gracefully
In some occasions KCL library calls shutdown multiple times for the same shards KCLRecordProcessor. We have to make sure it doesn't fail because closed dependencies.
1 parent 9c4daca commit fbd6754

File tree

2 files changed

+38
-13
lines changed

2 files changed

+38
-13
lines changed

source/src/main/java/com/trustpilot/connector/dynamodb/kcl/KclRecordProcessor.java

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -208,18 +208,19 @@ public void shutdown(ShutdownInput shutdownInput) {
208208
* Called on ShutdownReason.TERMINATE
209209
*/
210210
private void onTerminate(ShutdownInput shutdownInput) throws InvalidStateException, ShutdownException, InterruptedException {
211-
if (!lastProcessedSeqNo.isEmpty()) {
212-
ShardInfo processorRegister = shardRegister.get(this.shardId);
213-
214-
while (!processorRegister.getLastCommittedRecordSeqNo().equals(this.lastProcessedSeqNo)) {
215-
LOGGER.info(
216-
"Shard ended. Waiting for all data table: {} from shard: {} to be committed. " +
217-
"lastCommittedRecordSeqNo: {} lastProcessedSeqNo: {}",
218-
tableName,
219-
shardId,
220-
processorRegister.getLastCommittedRecordSeqNo(),
221-
this.lastProcessedSeqNo);
222-
Thread.sleep(1000);
211+
if (lastProcessedSeqNo != null && !lastProcessedSeqNo.isEmpty()) {
212+
ShardInfo processorRegister = shardRegister.getOrDefault(this.shardId, null);
213+
if (processorRegister != null) {
214+
while (!processorRegister.getLastCommittedRecordSeqNo().equals(this.lastProcessedSeqNo)) {
215+
LOGGER.info(
216+
"Shard ended. Waiting for all data table: {} from shard: {} to be committed. " +
217+
"lastCommittedRecordSeqNo: {} lastProcessedSeqNo: {}",
218+
tableName,
219+
shardId,
220+
processorRegister.getLastCommittedRecordSeqNo(),
221+
this.lastProcessedSeqNo);
222+
Thread.sleep(500);
223+
}
223224
}
224225
}
225226

@@ -231,7 +232,10 @@ private void onTerminate(ShutdownInput shutdownInput) throws InvalidStateExcepti
231232
"Shard ended. All data committed. Checkpoint and proceed to next one. Table: {} ShardID: {}",
232233
tableName,
233234
shardId);
234-
shutdownInput.getCheckpointer().checkpoint();
235+
IRecordProcessorCheckpointer checkpointer = shutdownInput.getCheckpointer();
236+
if (checkpointer != null) {
237+
checkpointer.checkpoint();
238+
}
235239
}
236240

237241
/**

source/src/test/java/com/trustpilot/connector/dynamodb/kcl/KclRecordProcessorTests.java

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -264,6 +264,27 @@ void shutdownWaitsForLastRecordToBeCommittedOnShardEndBeforeCheckpoint() throws
264264
verify(checkpointer, only()).checkpoint();
265265
}
266266

267+
@Test
268+
void shutdownSucceddsIfCalledMultipleTimes() throws InvalidStateException, ShutdownException {
269+
// Arrange
270+
IRecordProcessorCheckpointer checkpointer = Mockito.mock(IRecordProcessorCheckpointer.class);
271+
ProcessRecordsInput processRecordsInput = getProcessRecordsInput("SQ1").withCheckpointer(checkpointer);
272+
shardRegister.get(shardId).setLastCommittedRecordSeqNo("SQ1");
273+
274+
// Act
275+
processor.processRecords(processRecordsInput);
276+
ShutdownInput shutdownInput = new ShutdownInput()
277+
.withShutdownReason(ShutdownReason.TERMINATE)
278+
.withCheckpointer(checkpointer);
279+
280+
processor.shutdown(shutdownInput);
281+
processor.shutdown(shutdownInput);
282+
283+
// Assert
284+
verify(checkpointer, times(2)).checkpoint();
285+
assertFalse(shardRegister.containsKey(shardId));
286+
}
287+
267288
@Test
268289
void shutdownCheckpointsIfNoDataWasReceivedFromThisShard() throws InvalidStateException, ShutdownException {
269290
// Arrange

0 commit comments

Comments
 (0)