Skip to content

Commit a10c0c7

Browse files
graytaylor0graytaylor0
committed
Fix how failures are handled for shard partitions in DDB source
Signed-off-by: Taylor Gray <tylgry@amazon.com>
1 parent a875c19 commit a10c0c7

5 files changed

Lines changed: 148 additions & 11 deletions

File tree

‎data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardAcknowledgementManager.java‎

Lines changed: 31 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator;
1616
import org.opensearch.dataprepper.common.concurrent.BackgroundThreadFactory;
1717
import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourcePartition;
18+
import org.opensearch.dataprepper.model.source.coordinator.exceptions.PartitionUpdateException;
1819
import org.opensearch.dataprepper.plugins.source.dynamodb.DynamoDBSourceConfig;
1920
import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.StreamPartition;
2021
import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.state.StreamProgressState;
@@ -39,14 +40,16 @@
3940

4041
import org.opensearch.dataprepper.plugins.source.dynamodb.model.ShardCheckpointStatus;
4142

43+
import static org.opensearch.dataprepper.logging.DataPrepperMarkers.NOISY;
44+
4245
public class ShardAcknowledgementManager {
4346
private static final Logger LOG = LoggerFactory.getLogger(ShardAcknowledgementManager.class);
4447

4548
private static final String NULL_SEQUENCE_NUMBER = "null";
4649

4750
private static final long WAIT_FOR_ACKNOWLEDGMENTS_TIMEOUT = 10L;
4851

49-
static final Duration CHECKPOINT_INTERVAL = Duration.ofMinutes(2);
52+
static final Duration CHECKPOINT_INTERVAL = Duration.ofMinutes(3);
5053

5154
private final DynamoDBSourceConfig dynamoDBSourceConfig;
5255
private final Map<StreamPartition, ConcurrentLinkedQueue<ShardCheckpointStatus>> checkpoints = new ConcurrentHashMap<>();
@@ -153,13 +156,10 @@ boolean runMonitorAcknowledgmentLoop(final Consumer<StreamPartition> stopWorkerC
153156
sourceCoordinator.saveProgressStateForPartition(streamPartition, dynamoDBSourceConfig.getShardAcknowledgmentTimeout());
154157
LOG.debug("Checkpointed shard {} with latest sequence number acknowledged {}", streamPartition.getShardId(), latestCheckpointForShard.getSequenceNumber());
155158
}
156-
if (partitionsToGiveUp.contains(streamPartition)) {
157-
partitionsToRemove.add(streamPartition);
158-
sourceCoordinator.giveUpPartition(streamPartition);
159-
}
160-
161159
} catch (final Exception e) {
162-
LOG.error("Received exception while monitoring acknowledgments for stream partition {}", streamPartition.getPartitionKey(), e);
160+
LOG.error(NOISY, "Received exception while monitoring acknowledgments for stream partition {}, stop processing shard", streamPartition.getPartitionKey(), e);
161+
partitionsToRemove.add(streamPartition);
162+
partitionsToGiveUp.add(streamPartition);
163163
}
164164
}
165165

@@ -172,6 +172,13 @@ public AcknowledgementSet createAcknowledgmentSet(
172172
final boolean isFinalSetForPartition) {
173173
final String sequenceNumberNoNull = sequenceNumber == null ? NULL_SEQUENCE_NUMBER : sequenceNumber;
174174
final ShardCheckpointStatus shardCheckpointStatus = new ShardCheckpointStatus(sequenceNumber, Instant.now().toEpochMilli(), isFinalSetForPartition);
175+
176+
// Shard should already be in checkpoints map from call to startUpdatingOwnershipForShard, if it is not in the map
177+
// that means that ShardAcknowledgmentManager stopped tracking it due to some error, and another worker will pick it up
178+
// We throw an error in this case to have the ShardConsumer exit and stop reading data from the shard
179+
if (!isStillTrackingShard(streamPartition)) {
180+
throw new ShardNotTrackedException("The shard {} is not being tracked anymore, stop reading from shard");
181+
}
175182
checkpoints.computeIfAbsent(streamPartition, segment -> new ConcurrentLinkedQueue<>()).add(shardCheckpointStatus);
176183
ackStatuses.computeIfAbsent(streamPartition, segment -> new ConcurrentHashMap<>());
177184
ackStatuses.get(streamPartition).put(sequenceNumberNoNull, shardCheckpointStatus);
@@ -199,7 +206,13 @@ void updateOwnershipForAllShardPartitions() {
199206
if (Duration.between(lastCheckpointTime, Instant.now()).compareTo(CHECKPOINT_INTERVAL) > 0) {
200207
for (final StreamPartition streamPartition : checkpoints.keySet()) {
201208
if (!partitionsToRemove.contains(streamPartition)) {
202-
sourceCoordinator.saveProgressStateForPartition(streamPartition, dynamoDBSourceConfig.getShardAcknowledgmentTimeout());
209+
try {
210+
sourceCoordinator.saveProgressStateForPartition(streamPartition, dynamoDBSourceConfig.getShardAcknowledgmentTimeout());
211+
} catch (final PartitionUpdateException e) {
212+
LOG.warn(NOISY, "Failed to update progress state for shard {}, will stop tracking this shard as someone else owns it", streamPartition.getShardId());
213+
partitionsToRemove.add(streamPartition);
214+
partitionsToGiveUp.add(streamPartition);
215+
}
203216
}
204217
}
205218

@@ -215,8 +228,7 @@ private void handleFailure(final StreamPartition streamPartition,
215228
sourceCoordinator.saveProgressStateForPartition(streamPartition, dynamoDBSourceConfig.getShardAcknowledgmentTimeout());
216229
}
217230
partitionsToRemove.add(streamPartition);
218-
sourceCoordinator.giveUpPartition(streamPartition);
219-
partitionsToGiveUp.remove(streamPartition);
231+
partitionsToGiveUp.add(streamPartition);
220232
}
221233

222234
private void handleCompletedShard(final StreamPartition streamPartition) {
@@ -245,13 +257,17 @@ private void removePartitions() {
245257
ackStatuses.remove(streamPartition);
246258
});
247259

260+
partitionsToGiveUp.forEach(sourceCoordinator::giveUpPartition);
261+
248262
partitionsToRemove.clear();
263+
partitionsToGiveUp.clear();
249264
}
250265

251266
public void giveUpPartition(final StreamPartition streamPartition) {
252-
if (!partitionsToGiveUp.contains(streamPartition)) {
267+
if (isStillTrackingShard(streamPartition)) {
253268
LOG.debug("Adding partition {} to give up list", streamPartition.getPartitionKey());
254269
partitionsToGiveUp.add(streamPartition);
270+
partitionsToRemove.add(streamPartition);
255271
}
256272
}
257273

@@ -263,4 +279,8 @@ public boolean isExportDone(StreamPartition streamPartition) {
263279
public void startUpdatingOwnershipForShard(final StreamPartition streamPartition) {
264280
checkpoints.computeIfAbsent(streamPartition, segment -> new ConcurrentLinkedQueue<>());
265281
}
282+
283+
public boolean isStillTrackingShard(final StreamPartition streamPartition) {
284+
return !partitionsToRemove.contains(streamPartition) && checkpoints.containsKey(streamPartition);
285+
}
266286
}

‎data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -272,6 +272,11 @@ public void run() {
272272
lastCheckpointTime = System.currentTimeMillis();
273273
}
274274

275+
if (shardAcknowledgementManager != null && !shardAcknowledgementManager.isStillTrackingShard(streamPartition)) {
276+
LOG.warn("Shard {} is no longer being tracked by the acknowledgment manager, exiting", streamPartition.getShardId());
277+
break;
278+
}
279+
275280
GetRecordsResponse response = callGetRecords(shardIterator);
276281
shardIterator = response.nextShardIterator();
277282
if (!response.records().isEmpty()) {

‎data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardNotTrackedException.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
package org.opensearch.dataprepper.plugins.source.dynamodb.stream;
2+
3+
public class ShardNotTrackedException extends RuntimeException {
4+
public ShardNotTrackedException(String message) {
5+
super(message);
6+
}
7+
}

‎data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardAcknowledgementManagerTest.java‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
1818
import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager;
1919
import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator;
20+
import org.opensearch.dataprepper.model.source.coordinator.exceptions.PartitionUpdateException;
2021
import org.opensearch.dataprepper.plugins.source.dynamodb.DynamoDBSourceConfig;
2122
import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.StreamPartition;
2223
import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.state.StreamProgressState;
@@ -25,6 +26,11 @@
2526
import java.util.Optional;
2627
import java.util.function.Consumer;
2728

29+
import static org.hamcrest.MatcherAssert.assertThat;
30+
import static org.hamcrest.Matchers.equalTo;
31+
import static org.junit.jupiter.api.Assertions.assertThrows;
32+
import static org.mockito.Mockito.doThrow;
33+
import static org.mockito.Mockito.mock;
2834
import static org.mockito.Mockito.verifyNoInteractions;
2935
import static org.mockito.Mockito.verifyNoMoreInteractions;
3036

@@ -78,6 +84,8 @@ void testCreateAcknowledgmentSet() {
7884
when(acknowledgementSetManager.create(any(Consumer.class), any(Duration.class)))
7985
.thenReturn(acknowledgementSet);
8086

87+
shardAcknowledgementManager.startUpdatingOwnershipForShard(streamPartition);
88+
8189
AcknowledgementSet result = shardAcknowledgementManager.createAcknowledgmentSet(
8290
streamPartition, "seq123", false);
8391

@@ -91,6 +99,8 @@ void testCreateAcknowledgmentSet_withNullSequenceNumber() {
9199
when(acknowledgementSetManager.create(any(Consumer.class), any(Duration.class)))
92100
.thenReturn(acknowledgementSet);
93101

102+
shardAcknowledgementManager.startUpdatingOwnershipForShard(streamPartition);
103+
94104
AcknowledgementSet result = shardAcknowledgementManager.createAcknowledgmentSet(
95105
streamPartition, null, false);
96106

@@ -181,6 +191,8 @@ void testPositiveAcknowledgment_finalAcknowledgmentSet_completes_partition() {
181191
shardAcknowledgementManager.runMonitorAcknowledgmentLoop(stopWorkerConsumer);
182192

183193
verify(sourceCoordinator).completePartition(streamPartition);
194+
195+
assertThat(shardAcknowledgementManager.isStillTrackingShard(streamPartition), equalTo(false));
184196
}
185197

186198
@Test
@@ -199,6 +211,9 @@ void testAcknowledgmentTimeout() throws Exception {
199211
Thread.sleep(10);
200212

201213
// Run the acknowledgment loop
214+
assertThat(shardAcknowledgementManager.isStillTrackingShard(streamPartition), equalTo(true));
215+
shardAcknowledgementManager.runMonitorAcknowledgmentLoop(stopWorkerConsumer);
216+
assertThat(shardAcknowledgementManager.isStillTrackingShard(streamPartition), equalTo(false));
202217
shardAcknowledgementManager.runMonitorAcknowledgmentLoop(stopWorkerConsumer);
203218

204219
verify(sourceCoordinator).giveUpPartition(streamPartition);
@@ -215,6 +230,7 @@ void testMultipleAcknowledgments_handledCorrectly() {
215230
.thenReturn(acknowledgementSet);
216231
when(streamPartition.getProgressState()).thenReturn(Optional.of(streamProgressState));
217232

233+
shardAcknowledgementManager.startUpdatingOwnershipForShard(streamPartition);
218234
// Create first acknowledgment set and capture its callback
219235
shardAcknowledgementManager.createAcknowledgmentSet(streamPartition, "seq123", false);
220236
Consumer<Boolean> callback1 = ackCallbackCaptor.getValue();
@@ -228,6 +244,9 @@ void testMultipleAcknowledgments_handledCorrectly() {
228244
callback2.accept(false); // Negative ack
229245

230246
// Run the acknowledgment loop
247+
assertThat(shardAcknowledgementManager.isStillTrackingShard(streamPartition), equalTo(true));
248+
shardAcknowledgementManager.runMonitorAcknowledgmentLoop(stopWorkerConsumer);
249+
assertThat(shardAcknowledgementManager.isStillTrackingShard(streamPartition), equalTo(false));
231250
shardAcknowledgementManager.runMonitorAcknowledgmentLoop(stopWorkerConsumer);
232251

233252
final InOrder inOrder = Mockito.inOrder(sourceCoordinator, streamProgressState);
@@ -237,4 +256,67 @@ void testMultipleAcknowledgments_handledCorrectly() {
237256

238257
verifyNoMoreInteractions(streamProgressState);
239258
}
259+
260+
@Test
261+
void create_acknowledgment_set_throws_ShardNotTrackedException_if_shard_is_not_being_tracked() {
262+
assertThrows(ShardNotTrackedException.class, () -> shardAcknowledgementManager.createAcknowledgmentSet(streamPartition, "seq123", false));
263+
}
264+
265+
@Test
266+
void failure_to_update_state_for_one_partition_stops_tracking_that_partition_and_continues_on_others() {
267+
final StreamPartition secondPartition = mock(StreamPartition.class);
268+
269+
270+
when(streamPartition.getProgressState()).thenReturn(Optional.of(streamProgressState));
271+
272+
when(secondPartition.getProgressState()).thenReturn(Optional.of(streamProgressState));
273+
274+
// Setup
275+
when(dynamoDBSourceConfig.getShardAcknowledgmentTimeout()).thenReturn(Duration.ofMinutes(15));
276+
when(acknowledgementSetManager.create(ackCallbackCaptor.capture(), any(Duration.class)))
277+
.thenReturn(acknowledgementSet);
278+
when(streamPartition.getProgressState()).thenReturn(Optional.of(streamProgressState));
279+
280+
shardAcknowledgementManager.startUpdatingOwnershipForShard(streamPartition);
281+
// Create first acknowledgment set and capture its callback
282+
shardAcknowledgementManager.createAcknowledgmentSet(streamPartition, "seq123", false);
283+
Consumer<Boolean> callback1 = ackCallbackCaptor.getValue();
284+
285+
// Create second acknowledgment set and capture its callback
286+
shardAcknowledgementManager.createAcknowledgmentSet(streamPartition, "seq124", false);
287+
Consumer<Boolean> callback2 = ackCallbackCaptor.getValue();
288+
289+
shardAcknowledgementManager.startUpdatingOwnershipForShard(secondPartition);
290+
// Create first acknowledgment set and capture its callback
291+
shardAcknowledgementManager.createAcknowledgmentSet(secondPartition, "seq123", false);
292+
Consumer<Boolean> callback1_2 = ackCallbackCaptor.getValue();
293+
294+
// Create second acknowledgment set and capture its callback
295+
shardAcknowledgementManager.createAcknowledgmentSet(secondPartition, "seq124", false);
296+
Consumer<Boolean> callback2_2 = ackCallbackCaptor.getValue();
297+
298+
// Simulate mixed acknowledgments
299+
callback1.accept(true); // Positive ack
300+
callback2.accept(true);
301+
302+
callback1_2.accept(true);
303+
callback2_2.accept(true);
304+
305+
doThrow(PartitionUpdateException.class).when(sourceCoordinator).saveProgressStateForPartition(eq(streamPartition),
306+
any(Duration.class));
307+
308+
309+
// Run the acknowledgment loop
310+
assertThat(shardAcknowledgementManager.isStillTrackingShard(streamPartition), equalTo(true));
311+
assertThat(shardAcknowledgementManager.isStillTrackingShard(secondPartition), equalTo(true));
312+
shardAcknowledgementManager.runMonitorAcknowledgmentLoop(stopWorkerConsumer);
313+
assertThat(shardAcknowledgementManager.isStillTrackingShard(streamPartition), equalTo(false));
314+
assertThat(shardAcknowledgementManager.isStillTrackingShard(secondPartition), equalTo(true));
315+
shardAcknowledgementManager.runMonitorAcknowledgmentLoop(stopWorkerConsumer);
316+
317+
final InOrder inOrder = Mockito.inOrder(sourceCoordinator, streamProgressState);
318+
inOrder.verify(streamProgressState).setSequenceNumber("seq124");
319+
inOrder.verify(sourceCoordinator).saveProgressStateForPartition(secondPartition, dynamoDBSourceConfig.getShardAcknowledgmentTimeout());
320+
inOrder.verify(sourceCoordinator).giveUpPartition(streamPartition);
321+
}
240322
}

‎data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerTest.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@
5656
import static org.mockito.Mockito.mockStatic;
5757
import static org.mockito.Mockito.times;
5858
import static org.mockito.Mockito.verify;
59+
import static org.mockito.Mockito.verifyNoMoreInteractions;
5960
import static org.mockito.Mockito.when;
6061
import static org.opensearch.dataprepper.plugins.source.dynamodb.stream.ShardConsumer.BUFFER_TIMEOUT;
6162
import static org.opensearch.dataprepper.plugins.source.dynamodb.stream.ShardConsumer.DEFAULT_BUFFER_BATCH_SIZE;
@@ -173,6 +174,7 @@ void setup() throws Exception {
173174

174175

175176
lenient().when(shardAcknowledgementManager.isExportDone(any(StreamPartition.class))).thenReturn(true);
177+
lenient().when(shardAcknowledgementManager.isStillTrackingShard(any(StreamPartition.class))).thenReturn(true);
176178
}
177179

178180
@Test
@@ -369,6 +371,27 @@ void test_run_shardConsumer_calls_startUpdatingOwnershipForShard() throws Except
369371
verify(shardAcknowledgementManager).startUpdatingOwnershipForShard(streamPartition);
370372
}
371373

374+
@Test
375+
void test_run_shardConsumer_exits_if_acknowledgment_manager_stops_tracking_shard() throws Exception {
376+
when(shardAcknowledgementManager.isStillTrackingShard(streamPartition)).thenReturn(false);
377+
try (final MockedStatic<BufferAccumulator> bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class)) {
378+
bufferAccumulatorMockedStatic.when(() -> BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT)).thenReturn(bufferAccumulator);
379+
ShardConsumer shardConsumer = ShardConsumer.builder(dynamoDbStreamsClient, pluginMetrics, aggregateMetrics, buffer, streamConfig)
380+
.shardIterator(shardIterator)
381+
.shardAcknowledgementManager(shardAcknowledgementManager)
382+
.streamPartition(streamPartition)
383+
.tableInfo(tableInfo)
384+
.startTime(null)
385+
.waitForExport(false)
386+
.build();
387+
388+
shardConsumer.run();
389+
}
390+
391+
verify(shardAcknowledgementManager).startUpdatingOwnershipForShard(streamPartition);
392+
verifyNoMoreInteractions(shardAcknowledgementManager, dynamoDbStreamsClient);
393+
}
394+
372395
@Test
373396
void test_shard_has_records_null_iterator() throws Exception {
374397
final AcknowledgementSet finalAcknowledgementSet = mock(AcknowledgementSet.class);

0 commit comments

Comments
 (0)