This is an automated email from the ASF dual-hosted git repository.
freeznet pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pulsar-connectors.git
The following commit(s) were added to refs/heads/master by this push:
new 03a4a30a feat(kinesis): add message key mode for Kinesis records (#35)
03a4a30a is described below
commit 03a4a30a1fce544a73a77b437b4d693ce39e4c19
Author: Rui Fu <[email protected]>
AuthorDate: Sun Jul 12 10:05:16 2026 +0800
feat(kinesis): add message key mode for Kinesis records (#35)
---
.../apache/pulsar/io/kinesis/KinesisRecord.java | 12 +-
.../pulsar/io/kinesis/KinesisRecordProcessor.java | 144 +++++++++++++++----
.../pulsar/io/kinesis/KinesisSourceConfig.java | 26 ++++
.../io/kinesis/KinesisRecordProcessorTest.java | 157 ++++++++++++++++++---
.../pulsar/io/kinesis/KinesisRecordTest.java | 39 +++++
.../pulsar/io/kinesis/KinesisSourceConfigTest.java | 25 ++++
6 files changed, 359 insertions(+), 44 deletions(-)
diff --git
a/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecord.java
b/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecord.java
index 83530175..730bdfc8 100644
--- a/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecord.java
+++ b/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecord.java
@@ -43,9 +43,19 @@ public class KinesisRecord implements Record<byte[]> {
private final long subSequenceNumber;
private final KinesisRecordProcessor recordProcessor;
+ /** Convenience overload that defaults to {@link
KinesisSourceConfig.MessageKeyMode#PARTITION_KEY}. */
public KinesisRecord(KinesisClientRecord record, String shardId, long
millisBehindLatest,
Set<String> propertiesToInclude,
KinesisRecordProcessor recordProcessor) {
- this.key = Optional.of(record.partitionKey());
+ this(record, shardId, millisBehindLatest, propertiesToInclude,
recordProcessor,
+ KinesisSourceConfig.MessageKeyMode.PARTITION_KEY);
+ }
+
+ public KinesisRecord(KinesisClientRecord record, String shardId, long
millisBehindLatest,
+ Set<String> propertiesToInclude,
KinesisRecordProcessor recordProcessor,
+ KinesisSourceConfig.MessageKeyMode messageKeyMode) {
+ this.key = (messageKeyMode ==
KinesisSourceConfig.MessageKeyMode.SHARD_ID)
+ ? Optional.of(shardId)
+ : Optional.of(record.partitionKey());
this.sequenceNumber = record.sequenceNumber();
this.subSequenceNumber = record.subSequenceNumber();
this.recordProcessor = recordProcessor;
diff --git
a/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecordProcessor.java
b/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecordProcessor.java
index daee5041..c0879786 100644
---
a/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecordProcessor.java
+++
b/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecordProcessor.java
@@ -18,8 +18,11 @@
*/
package org.apache.pulsar.io.kinesis;
+import java.util.ArrayDeque;
+import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
@@ -50,13 +53,25 @@ public class KinesisRecordProcessor implements
ShardRecordProcessor {
private final LinkedBlockingQueue<KinesisRecord> queue;
private final SourceContext sourceContext;
private final Set<String> propertiesToInclude;
+ private final KinesisSourceConfig.MessageKeyMode messageKeyMode;
private final ScheduledExecutorService checkpointExecutor;
private final AtomicReference<RecordProcessorCheckpointer> checkpointerRef
= new AtomicReference<>();
private final AtomicBoolean isCheckpointing = new AtomicBoolean(false);
private String kinesisShardId;
private final AtomicInteger numRecordsInFlight = new AtomicInteger(0);
- private volatile CheckpointSequenceNumber sequenceNumberNeedToCheckpoint =
null;
+ // All checkpoint-position state below is guarded by checkpointLock. Acks
arrive concurrently on the
+ // downstream sink's send-completion callback threads, processRecords runs
on the KCL thread, and
+ // triggerCheckpoint runs on the checkpointExecutor thread, so the
deque/set/highest-acked triple must
+ // be mutated atomically as a group.
+ private final Object checkpointLock = new Object();
+ // Positions appended in processRecords, in KCL delivery (sequence) order.
+ private final ArrayDeque<CheckpointSequenceNumber> deliveredInOrder = new
ArrayDeque<>();
+ // Positions acked but not yet collapsed into the contiguous prefix (acked
ahead of an earlier gap).
+ private final HashSet<CheckpointSequenceNumber> ackedAwaitingCollapse =
new HashSet<>();
+ // The highest position whose entire delivery-order prefix has been acked.
The ONLY thing we checkpoint.
+ private CheckpointSequenceNumber highestContiguousAcked = null;
+
private volatile CheckpointSequenceNumber lastCheckpointSequenceNumber =
null;
public KinesisRecordProcessor(LinkedBlockingQueue<KinesisRecord> queue,
KinesisSourceConfig config,
@@ -66,6 +81,7 @@ public class KinesisRecordProcessor implements
ShardRecordProcessor {
this.numRetries = config.getNumRetries();
this.backoffTime = config.getBackoffTime();
this.propertiesToInclude = config.getPropertiesToInclude();
+ this.messageKeyMode = config.getMessageKeyMode();
this.sourceContext = sourceContext;
this.checkpointExecutor = checkpointExecutor;
}
@@ -78,36 +94,76 @@ public class KinesisRecordProcessor implements
ShardRecordProcessor {
checkpointer.checkpoint(checkpoint.sequenceNumber(),
checkpoint.subSequenceNumber());
lastCheckpointSequenceNumber = checkpoint;
log.info("Successfully checkpointed shard {} at {}",
kinesisShardId, checkpoint);
- isCheckpointing.set(false);
- checkpointExecutor.schedule(this::triggerCheckpoint,
checkpointInterval, TimeUnit.MILLISECONDS);
+ scheduleNextCheckpoint();
} catch (ThrottlingException | KinesisClientLibDependencyException e) {
if (attempt >= numRetries) {
- log.error("Checkpoint for shard {} failed after {} attempts at
{}. Terminating.",
- kinesisShardId, numRetries, checkpoint, e);
- sourceContext.fatal(e);
+ // A failed periodic checkpoint is recoverable: skip this
round instead of terminating the
+ // connector, and let the next scheduled checkpoint retry with
the latest acked position.
+ log.error("Checkpoint for shard {} failed after {} attempts at
{}. Skipping this round; "
+ + "will retry on the next interval.", kinesisShardId,
numRetries, checkpoint, e);
+ scheduleNextCheckpoint();
} else {
log.warn("Throttling/Dependency error on checkpoint for shard
{} at {}. Scheduling retry {} "
+ "after {}ms.", kinesisShardId, checkpoint,
attempt + 1, backoffTime);
- checkpointExecutor.schedule(() ->
tryCheckpointWithRetry(checkpointer, checkpoint, attempt + 1),
- backoffTime, TimeUnit.MILLISECONDS);
+ safeSchedule(() -> tryCheckpointWithRetry(checkpointer,
checkpoint, attempt + 1), backoffTime);
}
+ } catch (IllegalArgumentException e) {
+ // KCL rejects a backward / out-of-range checkpoint with
IllegalArgumentException. With the
+ // contiguous-prefix logic this should no longer happen, so
surface it as a metric for alerting,
+ // but still skip-and-reschedule rather than crashing the
connector.
+ log.warn("KCL rejected a backward/out-of-range checkpoint for
shard {} at {}. Skipping this round; "
+ + "will retry on the next interval.", kinesisShardId,
checkpoint, e);
+ sourceContext.recordMetric("kinesisBackwardCheckpointDetected", 1);
+ scheduleNextCheckpoint();
} catch (Exception e) {
- log.error("Caught a non-retryable exception for shard {} during
checkpoint at {}. Terminating.",
- kinesisShardId, checkpoint, e);
- sourceContext.fatal(e);
+ // Other non-retryable errors must NOT crash the connector. Skip
this round; a later ack with a
+ // higher position will advance the checkpoint on a subsequent
cycle.
+ log.error("Non-retryable exception during checkpoint for shard {}
at {}. Skipping this round; "
+ + "will retry on the next interval.", kinesisShardId,
checkpoint, e);
+ scheduleNextCheckpoint();
+ }
+ }
+
+ private void scheduleNextCheckpoint() {
+ isCheckpointing.set(false);
+ safeSchedule(this::triggerCheckpoint, checkpointInterval);
+ }
+
+ private void safeSchedule(Runnable task, long delayMillis) {
+ try {
+ checkpointExecutor.schedule(task, delayMillis,
TimeUnit.MILLISECONDS);
+ } catch (RejectedExecutionException e) {
+ // The checkpoint executor has been shut down (connector is
closing); nothing left to schedule.
+ log.debug("Checkpoint executor rejected a task for shard {}
(likely shutting down).", kinesisShardId);
}
}
public void updateSequenceNumberToCheckpoint(String sequenceNumber, long
subSequenceNumber) {
- CheckpointSequenceNumber newCheckpoint = new
CheckpointSequenceNumber(sequenceNumber, subSequenceNumber);
- log.debug("{} Updating sequence number to checkpoint {}",
kinesisShardId, newCheckpoint);
- this.sequenceNumberNeedToCheckpoint = newCheckpoint;
+ CheckpointSequenceNumber acked = new
CheckpointSequenceNumber(sequenceNumber, subSequenceNumber);
+ log.debug("{} Record acked at {}", kinesisShardId, acked);
+ synchronized (checkpointLock) {
+ ackedAwaitingCollapse.add(acked);
+ // Collapse the longest contiguous, fully-acked prefix (in
delivery order) into highestContiguousAcked.
+ // Stop at the first un-acked head: we must never advance the
checkpoint past a still-in-flight
+ // lower sequence, otherwise a crash would resume AFTER it and
lose it.
+ while (!deliveredInOrder.isEmpty()) {
+ CheckpointSequenceNumber head = deliveredInOrder.peekFirst();
+ if (!ackedAwaitingCollapse.remove(head)) {
+ break;
+ }
+ deliveredInOrder.pollFirst();
+ highestContiguousAcked = head;
+ }
+ }
this.numRecordsInFlight.decrementAndGet();
}
public void failed() {
+ // Intentionally do NOT remove this record's position from
deliveredInOrder: a record that failed to
+ // be delivered downstream must keep blocking the contiguous-prefix
collapse so the checkpoint can
+ // never advance past it. fatal() terminates the connector, which
reprocesses from the last checkpoint.
numRecordsInFlight.decrementAndGet();
- sourceContext.fatal(new PulsarClientException("Failed to process
Kinesis records due send to pulsar topic"));
+ sourceContext.fatal(new PulsarClientException("Failed to send Kinesis
record to Pulsar topic"));
}
@Override
@@ -116,25 +172,37 @@ public class KinesisRecordProcessor implements
ShardRecordProcessor {
log.info("Initializing KinesisRecordProcessor for shard {},
extendedSequenceNumber: {}, pendingCheckSeq: {}",
kinesisShardId, initializationInput.extendedSequenceNumber(),
initializationInput.pendingCheckpointSequenceNumber());
- checkpointExecutor.schedule(this::triggerCheckpoint,
checkpointInterval, TimeUnit.MILLISECONDS);
+ safeSchedule(this::triggerCheckpoint, checkpointInterval);
}
private void triggerCheckpoint() {
try {
if (isCheckpointing.compareAndSet(false, true)) {
final RecordProcessorCheckpointer checkpointer =
checkpointerRef.get();
- final CheckpointSequenceNumber currentCheckpoint =
this.sequenceNumberNeedToCheckpoint;
+ final CheckpointSequenceNumber currentCheckpoint;
+ final int pendingGap;
+ synchronized (checkpointLock) {
+ currentCheckpoint = this.highestContiguousAcked;
+ pendingGap = deliveredInOrder.size();
+ }
+ // Emit metrics outside the lock. A rising
kinesisCheckpointPendingGap means acks are stalled
+ // behind a gap (a data-loss exposure under the old
last-writer-wins code; here it only delays
+ // the checkpoint from advancing).
+ sourceContext.recordMetric("kinesisRecordsInFlight",
numRecordsInFlight.get());
+ sourceContext.recordMetric("kinesisCheckpointPendingGap",
pendingGap);
if (checkpointer != null && currentCheckpoint != null &&
!currentCheckpoint.equals(
lastCheckpointSequenceNumber)) {
tryCheckpointWithRetry(checkpointer, currentCheckpoint, 1);
} else {
- isCheckpointing.set(false);
- checkpointExecutor.schedule(this::triggerCheckpoint,
checkpointInterval, TimeUnit.MILLISECONDS);
+ scheduleNextCheckpoint();
}
}
} catch (Throwable e) {
- log.error("Error while triggering checkpoint for shard {}.
Terminating.", kinesisShardId, e);
- sourceContext.fatal(e);
+ // Defensive: never let a checkpoint-trigger failure crash the
connector. Keep the loop alive
+ // so the next scheduled cycle can make progress.
+ log.error("Unexpected error while triggering checkpoint for shard
{}. Skipping this round.",
+ kinesisShardId, e);
+ scheduleNextCheckpoint();
}
}
@@ -147,12 +215,24 @@ public class KinesisRecordProcessor implements
ShardRecordProcessor {
for (KinesisClientRecord record : processRecordsInput.records()) {
log.debug("Add record with sequence number {}:{} to queue for
shard {}.",
record.sequenceNumber(), record.subSequenceNumber(),
kinesisShardId);
+ // Register the delivery-order position BEFORE the record can be
read/acked from the queue, so a
+ // fast ack can never race ahead of its own deliveredInOrder
entry. KCL delivers each position at
+ // most once to a given processor instance, so no de-duplication
is needed here. queue.put may
+ // block, so the lock is released before the put.
+ synchronized (checkpointLock) {
+ deliveredInOrder.addLast(
+ new CheckpointSequenceNumber(record.sequenceNumber(),
record.subSequenceNumber()));
+ }
try {
queue.put(new KinesisRecord(record, this.kinesisShardId,
millisBehindLatest,
- propertiesToInclude, this));
+ propertiesToInclude, this, messageKeyMode));
} catch (Exception e) {
+ // The record never made it onto the queue, so it will never
be read or acked. Its position
+ // stays in deliveredInOrder so the contiguous prefix cannot
advance past this gap (no data
+ // loss); fatal() terminates the connector, which then
reprocesses from the last checkpoint.
log.error("Unable to create and queue KinesisRecord for shard
{}.", kinesisShardId, e);
sourceContext.fatal(e);
+ return;
}
numRecordsInFlight.incrementAndGet();
}
@@ -169,7 +249,15 @@ public class KinesisRecordProcessor implements
ShardRecordProcessor {
numRecordsInFlight.get(), kinesisShardId);
try {
for (int i = 0; i < numRetries; i++) {
- if (numRecordsInFlight.get() == 0) {
+ // "Fully processed" means the entire delivery-order prefix
has been acked (deque empty), not
+ // merely that in-flight reached 0: a record that failed
downstream decrements in-flight but
+ // intentionally keeps its position in deliveredInOrder.
Gating SHARD_END on deque-empty keeps
+ // the full checkpoint() from ever marking the shard complete
while a gap remains.
+ boolean drained;
+ synchronized (checkpointLock) {
+ drained = deliveredInOrder.isEmpty();
+ }
+ if (drained) {
processedInTime = true;
break;
}
@@ -197,14 +285,18 @@ public class KinesisRecordProcessor implements
ShardRecordProcessor {
} else {
log.warn("Not all records for shard {} were processed or not a
shard end. "
+ "Performing best-effort checkpoint.",
kinesisShardId);
- final CheckpointSequenceNumber finalCheckpoint =
this.sequenceNumberNeedToCheckpoint;
+ final CheckpointSequenceNumber finalCheckpoint;
+ synchronized (checkpointLock) {
+ finalCheckpoint = this.highestContiguousAcked;
+ }
if (finalCheckpoint != null) {
checkpointer.checkpoint(finalCheckpoint.sequenceNumber(),
finalCheckpoint.subSequenceNumber());
}
}
} catch (Exception e) {
+ // Defensive: a failed final checkpoint must not crash the
connector. Under at-least-once
+ // semantics the un-checkpointed records will simply be
reprocessed by the next lease owner.
log.error("Failed to perform final checkpoint for shard {}. Data
may be reprocessed.", kinesisShardId, e);
- sourceContext.fatal(e);
}
}
@@ -219,4 +311,4 @@ public class KinesisRecordProcessor implements
ShardRecordProcessor {
log.info("Shutdown requested for shard {}, starting final checkpoint
process.", kinesisShardId);
finalizeAndCheckpoint(shutdownRequestedInput.checkpointer(), false);
}
-}
\ No newline at end of file
+}
diff --git
a/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisSourceConfig.java
b/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisSourceConfig.java
index ebcd06bc..01d17af1 100644
---
a/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisSourceConfig.java
+++
b/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisSourceConfig.java
@@ -144,6 +144,32 @@ public class KinesisSourceConfig extends BaseKinesisConfig
implements Serializab
+ "kinesis.partition.key,kinesis.sequence.number";
private transient Set<String> propertiesToInclude;
+ /**
+ * Determines which value is used as the Pulsar message key for each
Kinesis record.
+ */
+ public enum MessageKeyMode {
+ /**
+ * Use the Kinesis record partition key (default; preserves existing
downstream partitioning).
+ */
+ PARTITION_KEY,
+ /**
+ * Use the Kinesis shard ID as the message key, routing all records
from the same shard to
+ * the same downstream partition. This preserves per-shard ordering
and reduces checkpoint
+ * reordering gaps, but may create hot partitions if shard traffic is
skewed.
+ */
+ SHARD_ID
+ }
+
+ @FieldDoc(
+ required = false,
+ defaultValue = "PARTITION_KEY",
+ help = "Which value to use as the Pulsar message key for each record. "
+ + "PARTITION_KEY (default) uses the Kinesis record partition
key (preserves existing downstream "
+ + "partitioning). SHARD_ID routes all records of a Kinesis
shard to the same key, keeping them on "
+ + "one downstream partition to preserve per-shard ordering and
reduce checkpoint reordering gaps; "
+ + "note this changes downstream partition distribution and may
create hot partitions.")
+ private MessageKeyMode messageKeyMode = MessageKeyMode.PARTITION_KEY;
+
public static KinesisSourceConfig load(Map<String, Object> config,
SourceContext sourceContext) {
KinesisSourceConfig kinesisSourceConfig =
IOConfigUtils.loadWithSecrets(config,
KinesisSourceConfig.class, sourceContext);
diff --git
a/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisRecordProcessorTest.java
b/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisRecordProcessorTest.java
index 70c04fa8..52a4b805 100644
---
a/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisRecordProcessorTest.java
+++
b/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisRecordProcessorTest.java
@@ -41,6 +41,7 @@ import
software.amazon.kinesis.lifecycle.events.InitializationInput;
import software.amazon.kinesis.lifecycle.events.ProcessRecordsInput;
import software.amazon.kinesis.lifecycle.events.ShardEndedInput;
import software.amazon.kinesis.lifecycle.events.ShutdownRequestedInput;
+import software.amazon.kinesis.exceptions.ThrottlingException;
import software.amazon.kinesis.processor.RecordProcessorCheckpointer;
import software.amazon.kinesis.retrieval.KinesisClientRecord;
@@ -67,6 +68,7 @@ public class KinesisRecordProcessorTest {
when(config.getNumRetries()).thenReturn(1);
when(config.getBackoffTime()).thenReturn(100L);
when(config.getPropertiesToInclude()).thenReturn(Collections.emptySet());
+
when(config.getMessageKeyMode()).thenReturn(KinesisSourceConfig.MessageKeyMode.PARTITION_KEY);
recordProcessor = new KinesisRecordProcessor(queue, config,
sourceContext, checkpointExecutor);
}
@@ -93,32 +95,105 @@ public class KinesisRecordProcessorTest {
}
@Test
- public void testCheckpointLogicOnlyAdvancesForward() throws Exception {
- // Arrange
+ public void testOutOfOrderAcksDoNotAdvancePastGap() throws Exception {
+ // Arrange: deliver three records of the same shard IN ORDER via
processRecords. KCL delivery order
+ // is authoritative, so the contiguous prefix is seq-1/0 -> seq-1/1 ->
seq-1/2.
recordProcessor.initialize(createMockInitializationInput());
verify(checkpointExecutor, Mockito.times(1))
.schedule(scheduledTaskCaptor.capture(), anyLong(),
any(TimeUnit.class));
Runnable scheduledCheckpointTask = scheduledTaskCaptor.getValue();
- ProcessRecordsInput processRecordsInput =
createMockProcessRecordsInput(checkpointer);
- recordProcessor.processRecords(processRecordsInput);
+ recordProcessor.processRecords(createMockProcessRecordsInput(
+ createMockKinesisRecord("seq-1", 0L),
+ createMockKinesisRecord("seq-1", 1L),
+ createMockKinesisRecord("seq-1", 2L)));
- // Act & Assert 1: Ack a later sub-sequence number first.
- recordProcessor.updateSequenceNumberToCheckpoint("seq-1", 5L);
+ // Act & Assert 1: ack ONLY the highest (seq-1/2). There is still a
gap at seq-1/0, so the
+ // contiguous prefix is empty and nothing must be checkpointed.
+ recordProcessor.updateSequenceNumberToCheckpoint("seq-1", 2L);
+ scheduledCheckpointTask.run();
+ verify(checkpointer, never()).checkpoint(any(String.class), anyLong());
+
+ // Act & Assert 2: ack the lowest (seq-1/0). The contiguous prefix now
covers exactly seq-1/0,
+ // because seq-1/1 is still un-acked. Checkpoint must be seq-1/0, NOT
the already-acked seq-1/2.
+ recordProcessor.updateSequenceNumberToCheckpoint("seq-1", 0L);
+ scheduledCheckpointTask.run();
+ verify(checkpointer, times(1)).checkpoint("seq-1", 0L);
+
+ // Act & Assert 3: ack the middle (seq-1/1). The prefix now collapses
through seq-1/1 and the
+ // already-acked seq-1/2, advancing the checkpoint to seq-1/2.
+ recordProcessor.updateSequenceNumberToCheckpoint("seq-1", 1L);
+ scheduledCheckpointTask.run();
+ verify(checkpointer, times(1)).checkpoint("seq-1", 2L);
+ }
+
+ @Test
+ public void testOutOfOrderAcksAcrossSequenceNumbers() throws Exception {
+ // Distinct sequence numbers are delivered in order seq-100, seq-101,
seq-102.
+ recordProcessor.initialize(createMockInitializationInput());
+ verify(checkpointExecutor, Mockito.times(1))
+ .schedule(scheduledTaskCaptor.capture(), anyLong(),
any(TimeUnit.class));
+ Runnable scheduledCheckpointTask = scheduledTaskCaptor.getValue();
+ recordProcessor.processRecords(createMockProcessRecordsInput(
+ createMockKinesisRecord("seq-100", 0L),
+ createMockKinesisRecord("seq-101", 0L),
+ createMockKinesisRecord("seq-102", 0L)));
+
+ // Ack 102 then 100; 101 is still in-flight. The checkpoint must be
the contiguous head seq-100,
+ // never the higher seq-102 (which would skip the un-acked seq-101 and
lose it on resume).
+ recordProcessor.updateSequenceNumberToCheckpoint("seq-102", 0L);
+ recordProcessor.updateSequenceNumberToCheckpoint("seq-100", 0L);
+ scheduledCheckpointTask.run();
+ verify(checkpointer, times(1)).checkpoint("seq-100", 0L);
+ verify(checkpointer, never()).checkpoint("seq-102", 0L);
+
+ // Now ack 101: the prefix collapses through 101 and the already-acked
102, advancing to seq-102.
+ recordProcessor.updateSequenceNumberToCheckpoint("seq-101", 0L);
scheduledCheckpointTask.run();
- verify(checkpointer, times(1)).checkpoint("seq-1", 5L);
+ verify(checkpointer, times(1)).checkpoint("seq-102", 0L);
+ }
+
+ @Test
+ public void testNonRetryableCheckpointFailureDoesNotCrashConnector()
throws Exception {
+ // Arrange: initialize schedules the first checkpoint task.
+ recordProcessor.initialize(createMockInitializationInput());
+ verify(checkpointExecutor, times(1))
+ .schedule(scheduledTaskCaptor.capture(), anyLong(),
any(TimeUnit.class));
+ Runnable scheduledCheckpointTask = scheduledTaskCaptor.getValue();
+
recordProcessor.processRecords(createMockProcessRecordsInput(checkpointer));
+
+ // Simulate KCL rejecting a backward / out-of-range checkpoint (an
IllegalArgumentException).
+ Mockito.doThrow(new IllegalArgumentException("Could not checkpoint ...
did not fall into acceptable range"))
+ .when(checkpointer).checkpoint(any(String.class), anyLong());
+ recordProcessor.updateSequenceNumberToCheckpoint("seq-1", 5L);
- // Act & Assert 2: Ack an earlier (out-of-order) sub-sequence number.
- // The checkpointToCommit is now older than
lastSuccessfullyCheckpointed.
- recordProcessor.updateSequenceNumberToCheckpoint("seq-1", 3L);
+ // Act
scheduledCheckpointTask.run();
- // Verify checkpoint was NOT called again, because the position is not
new.
- verify(checkpointer, times(1)).checkpoint("seq-1", 5L);
- // Act & Assert 3: Ack a new, even later sub-sequence number.
- recordProcessor.updateSequenceNumberToCheckpoint("seq-1", 7L);
+ // Assert: the connector must not be terminated, and the checkpoint
loop must be rescheduled
+ // (one schedule from initialize + one from the failed round).
+ verify(sourceContext, never()).fatal(any());
+ verify(checkpointExecutor, times(2)).schedule(any(Runnable.class),
anyLong(), any(TimeUnit.class));
+ }
+
+ @Test
+ public void testExhaustedRetriableCheckpointFailureDoesNotCrashConnector()
throws Exception {
+ // Arrange: numRetries == 1 (see setup), so the first throttling
failure already exhausts retries.
+ recordProcessor.initialize(createMockInitializationInput());
+ verify(checkpointExecutor, times(1))
+ .schedule(scheduledTaskCaptor.capture(), anyLong(),
any(TimeUnit.class));
+ Runnable scheduledCheckpointTask = scheduledTaskCaptor.getValue();
+
recordProcessor.processRecords(createMockProcessRecordsInput(checkpointer));
+
+ Mockito.doThrow(new ThrottlingException("throttled"))
+ .when(checkpointer).checkpoint(any(String.class), anyLong());
+ recordProcessor.updateSequenceNumberToCheckpoint("seq-1", 1L);
+
+ // Act
scheduledCheckpointTask.run();
- // Verify checkpoint is now called with the new, advanced position.
- verify(checkpointer, times(1)).checkpoint("seq-1", 7L);
+
+ // Assert: retries exhausted, but the connector survives and the loop
is rescheduled.
+ verify(sourceContext, never()).fatal(any());
+ verify(checkpointExecutor, times(2)).schedule(any(Runnable.class),
anyLong(), any(TimeUnit.class));
}
@Test(timeOut = 3000)
@@ -160,6 +235,54 @@ public class KinesisRecordProcessorTest {
verify(shutdownCheckpointer, times(1)).checkpoint("seq-shutdown-1",
0L);
}
+ @Test(timeOut = 5000)
+ public void testShardEndedAllAckedPerformsFullCheckpoint() throws
Exception {
+ // Arrange: deliver two records and ack BOTH, so the contiguous prefix
fully drains (deque empty).
+ recordProcessor.processRecords(createMockProcessRecordsInput(
+ createMockKinesisRecord("seq-A", 0L),
+ createMockKinesisRecord("seq-B", 1L)
+ ));
+ queue.take().ack();
+ queue.take().ack();
+
+ RecordProcessorCheckpointer shardEndCheckpointer =
Mockito.mock(RecordProcessorCheckpointer.class);
+ ShardEndedInput shardEndedInput =
createMockShardEndedInput(shardEndCheckpointer);
+
+ // Act
+ recordProcessor.shardEnded(shardEndedInput);
+
+ // Assert: everything is acked, so KCL is told the shard is fully
consumed via the no-arg checkpoint().
+ verify(shardEndCheckpointer, times(1)).checkpoint();
+ verify(shardEndCheckpointer, never()).checkpoint(any(String.class),
anyLong());
+ }
+
+ @Test(timeOut = 5000)
+ public void testShardEndedWithGapDoesNotPerformFullCheckpoint() throws
Exception {
+ // Arrange: deliver three records; ack the 1st and 3rd, fail the 2nd.
in-flight reaches 0 but a gap
+ // (seq-B) remains in the delivery-order deque, so SHARD_END must NOT
mark the shard fully consumed.
+ recordProcessor.processRecords(createMockProcessRecordsInput(
+ createMockKinesisRecord("seq-A", 0L),
+ createMockKinesisRecord("seq-B", 1L),
+ createMockKinesisRecord("seq-C", 2L)
+ ));
+ KinesisRecord recA = queue.take();
+ KinesisRecord recB = queue.take();
+ KinesisRecord recC = queue.take();
+ recA.ack();
+ recC.ack(); // acked ahead of the gap; must not collapse past the
un-acked seq-B
+ recB.fail(); // sourceContext is a mock, so fatal() is a no-op and
the processor keeps running
+
+ RecordProcessorCheckpointer shardEndCheckpointer =
Mockito.mock(RecordProcessorCheckpointer.class);
+ ShardEndedInput shardEndedInput =
createMockShardEndedInput(shardEndCheckpointer);
+
+ // Act
+ recordProcessor.shardEnded(shardEndedInput);
+
+ // Assert: best-effort checkpoint at the contiguous prefix (seq-A)
only; no full SHARD_END checkpoint.
+ verify(shardEndCheckpointer, times(1)).checkpoint("seq-A", 0L);
+ verify(shardEndCheckpointer, never()).checkpoint();
+ }
+
private KinesisClientRecord createMockKinesisRecord(String sequenceNumber,
long subSequenceNumber) {
KinesisClientRecord mockRecord =
Mockito.mock(KinesisClientRecord.class);
when(mockRecord.partitionKey()).thenReturn("test-key");
@@ -198,4 +321,4 @@ public class KinesisRecordProcessorTest {
when(input.checkpointer()).thenReturn(checkpointer);
return input;
}
-}
\ No newline at end of file
+}
diff --git
a/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisRecordTest.java
b/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisRecordTest.java
index 6383cbab..29ba6bfd 100644
--- a/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisRecordTest.java
+++ b/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisRecordTest.java
@@ -110,4 +110,43 @@ public class KinesisRecordTest {
assertTrue(properties.isEmpty());
}
+
+ // --- MessageKeyMode tests ---
+
+ @Test
+ public void testDefaultKeyModeUsesPartitionKey() {
+ // The 5-arg constructor defaults to PARTITION_KEY.
+ KinesisRecord record = new KinesisRecord(mockRecord, shardId,
millisBehindLatest,
+ Collections.emptySet(), null);
+ assertTrue(record.getKey().isPresent());
+ assertEquals(record.getKey().get(), partitionKey);
+ }
+
+ @Test
+ public void testExplicitPartitionKeyMode() {
+ KinesisRecord record = new KinesisRecord(mockRecord, shardId,
millisBehindLatest,
+ Collections.emptySet(), null,
KinesisSourceConfig.MessageKeyMode.PARTITION_KEY);
+ assertTrue(record.getKey().isPresent());
+ assertEquals(record.getKey().get(), partitionKey);
+ }
+
+ @Test
+ public void testShardIdKeyMode() {
+ KinesisRecord record = new KinesisRecord(mockRecord, shardId,
millisBehindLatest,
+ Collections.emptySet(), null,
KinesisSourceConfig.MessageKeyMode.SHARD_ID);
+ assertTrue(record.getKey().isPresent());
+ assertEquals(record.getKey().get(), shardId);
+ }
+
+ @Test
+ public void testShardIdKeyModeDoesNotAffectPartitionKeyProperty() {
+ // Even in SHARD_ID mode, the kinesis.partition.key *property* must
still reflect the
+ // original record partition key (property semantics are independent
of the message key).
+ Set<String> propertiesToInclude =
Collections.singleton(KinesisRecord.PARTITION_KEY);
+ KinesisRecord record = new KinesisRecord(mockRecord, shardId,
millisBehindLatest,
+ propertiesToInclude, null,
KinesisSourceConfig.MessageKeyMode.SHARD_ID);
+ assertTrue(record.getKey().isPresent());
+ assertEquals(record.getKey().get(), shardId);
+ assertEquals(record.getProperties().get(KinesisRecord.PARTITION_KEY),
partitionKey);
+ }
}
\ No newline at end of file
diff --git
a/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisSourceConfigTest.java
b/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisSourceConfigTest.java
index 0cedb609..2db085d3 100644
---
a/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisSourceConfigTest.java
+++
b/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisSourceConfigTest.java
@@ -211,4 +211,29 @@ public class KinesisSourceConfigTest {
Set<String> properties = config.getPropertiesToInclude();
assertTrue(properties.isEmpty());
}
+
+ @Test
+ public final void messageKeyModeDefaultsToPartitionKey() {
+ Map<String, Object> map = new HashMap<>();
+ map.put("awsRegion", "us-west-1");
+ map.put("awsKinesisStreamName", "my-stream");
+ map.put("awsCredentialPluginParam",
"{\"accessKey\":\"myKey\",\"secretKey\":\"my-Secret\"}");
+
+ KinesisSourceConfig config = KinesisSourceConfig.load(map,
Mockito.mock(SourceContext.class));
+
+ assertEquals(config.getMessageKeyMode(),
KinesisSourceConfig.MessageKeyMode.PARTITION_KEY);
+ }
+
+ @Test
+ public final void messageKeyModeDeserializesShardIdFromString() {
+ Map<String, Object> map = new HashMap<>();
+ map.put("awsRegion", "us-west-1");
+ map.put("awsKinesisStreamName", "my-stream");
+ map.put("awsCredentialPluginParam",
"{\"accessKey\":\"myKey\",\"secretKey\":\"my-Secret\"}");
+ map.put("messageKeyMode", "SHARD_ID");
+
+ KinesisSourceConfig config = KinesisSourceConfig.load(map,
Mockito.mock(SourceContext.class));
+
+ assertEquals(config.getMessageKeyMode(),
KinesisSourceConfig.MessageKeyMode.SHARD_ID);
+ }
}
\ No newline at end of file
