This is an automated email from the ASF dual-hosted git repository.
lizhimins pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/rocketmq-flink.git
The following commit(s) were added to refs/heads/main by this push:
new 497fb05 [ISSUE #161] fix(source): keep failed-topic queues allocated
and resume re-added queues from checkpoint offsets (#162)
497fb05 is described below
commit 497fb05abeadb3cd9b413f1f5ab62394932730e1
Author: lizhimins <[email protected]>
AuthorDate: Fri Nov 6 18:57:16 2026 +0800
[ISSUE #161] fix(source): keep failed-topic queues allocated and resume
re-added queues from checkpoint offsets (#162)
Two gaps remained in periodic partition discovery after #152:
1. A partial route-fetch failure across multiple topics still revoked the
failed topic's queues: requestServiceDiscovery only skipped the round
when every topic failed, and the plain set difference against the
allocated set reported the failed topic's queues as decreased. The
discovery callable now returns the failed topics together with the
discovered queues, and the coordinator excludes their queues from the
decrease evaluation. A failed fetch is no evidence of absence.
2. A revoked queue that reappeared restarted from the configured starting
selector: initializeSourceSplits discarded the checkedOffsets entry on
decrease, so the re-added queue was re-resolved (with EARLIEST, replaying
the whole queue). The last checkpointed offset is now retained for
revoked queues and preferred when the queue reappears, the same rule the
migration path already applies to surviving queues. The selector still
applies to genuinely new queues, and after coordinator failover where
the retained offsets are lost.
RocketMQSourceEnumeratorDiscoveryToleranceTest now covers the per-topic
exclusion on partial failures at both the handler and the diff level; the
new RocketMQSourceEnumeratorRevokedOffsetResumeTest covers checkpoint
resume, new-queue selector fallback, mixed rounds, and the
NO_STOPPING_OFFSET
guard.
---
.../enumerator/RocketMQSourceEnumerator.java | 70 +++++-
...etMQSourceEnumeratorDiscoveryToleranceTest.java | 119 ++++++++--
...tMQSourceEnumeratorRevokedOffsetResumeTest.java | 241 +++++++++++++++++++++
3 files changed, 396 insertions(+), 34 deletions(-)
diff --git
a/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/source/enumerator/RocketMQSourceEnumerator.java
b/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/source/enumerator/RocketMQSourceEnumerator.java
index d801c5f..cad9884 100644
---
a/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/source/enumerator/RocketMQSourceEnumerator.java
+++
b/flink-connector-rocketmq/src/main/java/org/apache/flink/connector/rocketmq/source/enumerator/RocketMQSourceEnumerator.java
@@ -83,6 +83,13 @@ public class RocketMQSourceEnumerator
// Used for queue dynamic allocation
private final Map<MessageQueue, Long> checkedOffsets;
+
+ // Last checkpointed offsets of revoked queues, retained so a queue that
reappears (broker
+ // rolling restart, scale-down then scale-up) resumes from checkpointed
progress instead of
+ // the configured starting selector. Coordinator-local; lost on failover,
where the selector
+ // applies again.
+ private final Map<MessageQueue, Long> revokedQueueOffsets;
+
private boolean[] initTask;
// The internal states of the enumerator.
@@ -136,6 +143,7 @@ public class RocketMQSourceEnumerator
// Support allocate splits to reader
this.checkedOffsets = new ConcurrentHashMap<>();
+ this.revokedQueueOffsets = new ConcurrentHashMap<>();
this.reflectedQueueToTaskId = new ConcurrentHashMap<>();
this.pendingSplitAssignmentMap = new ConcurrentHashMap<>();
this.allocatedSet = new ConcurrentHashMap<>();
@@ -258,7 +266,7 @@ public class RocketMQSourceEnumerator
// ----------------- private methods -------------------
- private Set<MessageQueue> requestServiceDiscovery() {
+ private ServiceDiscoveryResult requestServiceDiscovery() {
// Ensure all subtasks have been initialized
try {
if (initTask != null) {
@@ -294,8 +302,9 @@ public class RocketMQSourceEnumerator
// A failed fetch says nothing about the current route: the topic
may be perfectly
// healthy behind an unreachable nameserver. Reporting its queues
as gone would
// revoke and later re-initialize their splits (mass rebalance
plus an offset
- // rewind once the nameserver is back), so failed topics are
excluded from this
- // round instead, and a round where every topic failed is skipped
entirely.
+ // rewind once the nameserver is back), so the coordinator
excludes their queues
+ // from the decrease evaluation instead, and a round where every
topic failed is
+ // skipped entirely.
log.warn(
"Skipping partition decrease evaluation for topics whose
route fetch failed: {}",
failedTopics);
@@ -303,17 +312,18 @@ public class RocketMQSourceEnumerator
return null;
}
}
- return fetchedQueues;
+ return new ServiceDiscoveryResult(fetchedQueues, failedTopics);
}
// This method should only be invoked in the coordinator executor thread.
- private void handleSourceQueueChange(Set<MessageQueue> latestSet,
Throwable t) {
+ private void handleSourceQueueChange(ServiceDiscoveryResult
discoveryResult, Throwable t) {
if (t != null) {
throw new FlinkRuntimeException("Failed to handle source splits
change due to ", t);
}
- if (latestSet == null) {
+ if (discoveryResult == null) {
return;
}
+ Set<MessageQueue> latestSet = discoveryResult.queues;
if (latestSet.isEmpty() && !allocatedSet.isEmpty()) {
// An empty discovery result while queues are allocated is
indistinguishable from
// "every route fetch failed" and almost never a real route
deletion; treating it
@@ -326,7 +336,8 @@ public class RocketMQSourceEnumerator
return;
}
- final SourceChangeResult sourceChangeResult =
getSourceChangeResult(latestSet);
+ final SourceChangeResult sourceChangeResult =
+ getSourceChangeResult(latestSet, discoveryResult.failedTopics);
if (sourceChangeResult.isEmpty()) {
log.debug("Skip handle source allocated due to not queue change");
return;
@@ -346,8 +357,22 @@ public class RocketMQSourceEnumerator
OffsetsSelector.MessageQueueOffsetsRetriever offsetsRetriever =
new InnerConsumerImpl.RemotingOffsetsRetrieverImpl(consumer);
- Map<MessageQueue, Long> increaseStartingOffsets =
- startingOffsetsSelector.getMessageQueueOffsets(increaseSet,
offsetsRetriever);
+ // A reappearing queue resumes from its last checkpointed offset, the
same rule the
+ // migration path applies to surviving queues; re-resolving with the
configured
+ // selector (e.g. EARLIEST) would replay the whole queue. Only
genuinely new queues
+ // go through the selector.
+ Map<MessageQueue, Long> increaseStartingOffsets = new HashMap<>();
+ Set<MessageQueue> unseenQueues = new HashSet<>();
+ for (MessageQueue mq : increaseSet) {
+ Long checkpointedOffset = revokedQueueOffsets.get(mq);
+ if (checkpointedOffset != null) {
+ increaseStartingOffsets.put(mq, checkpointedOffset);
+ } else {
+ unseenQueues.add(mq);
+ }
+ }
+ increaseStartingOffsets.putAll(
+ startingOffsetsSelector.getMessageQueueOffsets(unseenQueues,
offsetsRetriever));
Map<MessageQueue, Long> increaseStoppingOffsets =
stoppingOffsetsSelector.getMessageQueueOffsets(increaseSet,
offsetsRetriever);
Map<MessageQueue, Long> decreaseStoppingOffsets =
@@ -384,12 +409,19 @@ public class RocketMQSourceEnumerator
decreaseStoppingOffsets.getOrDefault(
mq,
RocketMQSourceSplit.NO_STOPPING_OFFSET);
allocatedSet.remove(mq);
- checkedOffsets.remove(mq);
+ Long lastCheckpointedOffset =
checkedOffsets.remove(mq);
+ if (lastCheckpointedOffset != null
+ && lastCheckpointedOffset >= 0) {
+ revokedQueueOffsets.put(mq,
lastCheckpointedOffset);
+ }
return new RocketMQSourceSplit(
mq, startingOffset,
stoppingOffset, false);
})
.collect(Collectors.toSet());
+ // The re-added queues are tracked in checkedOffsets again.
+ increaseSet.forEach(revokedQueueOffsets::remove);
+
return new SourceSplitChangeResult(increaseSplitSet, decreaseSplitSet);
}
@@ -686,6 +718,18 @@ public class RocketMQSourceEnumerator
}
}
+ /** The queues discovered in one round plus the topics whose route fetch
failed. */
+ @VisibleForTesting
+ static class ServiceDiscoveryResult {
+ final Set<MessageQueue> queues;
+ final Set<String> failedTopics;
+
+ ServiceDiscoveryResult(Set<MessageQueue> queues, Set<String>
failedTopics) {
+ this.queues = queues;
+ this.failedTopics = failedTopics;
+ }
+ }
+
/** A container class to hold the newly added partitions and removed
partitions. */
@VisibleForTesting
private static class SourceChangeResult {
@@ -737,12 +781,16 @@ public class RocketMQSourceEnumerator
}
@VisibleForTesting
- private SourceChangeResult getSourceChangeResult(Set<MessageQueue>
latestSet) {
+ private SourceChangeResult getSourceChangeResult(
+ Set<MessageQueue> latestSet, Set<String> failedTopics) {
Set<MessageQueue> currentSet =
Collections.unmodifiableSet(this.allocatedSet.keySet());
Set<MessageQueue> increaseSet = new HashSet<>(latestSet);
increaseSet.removeAll(currentSet);
Set<MessageQueue> decreaseSet = new HashSet<>(currentSet);
decreaseSet.removeAll(latestSet);
+ // A failed route fetch is no evidence that the topic's queues are
gone: keep them
+ // allocated until a successful fetch shows they are really absent.
+ decreaseSet.removeIf(mq -> failedTopics.contains(mq.getTopic()));
SourceChangeResult changeResult = new SourceChangeResult(increaseSet,
decreaseSet);
diff --git
a/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/source/enumerator/RocketMQSourceEnumeratorDiscoveryToleranceTest.java
b/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/source/enumerator/RocketMQSourceEnumeratorDiscoveryToleranceTest.java
index a35a614..a6e2aee 100644
---
a/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/source/enumerator/RocketMQSourceEnumeratorDiscoveryToleranceTest.java
+++
b/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/source/enumerator/RocketMQSourceEnumeratorDiscoveryToleranceTest.java
@@ -43,7 +43,8 @@ import static org.mockito.Mockito.when;
/**
* A route fetch failure must never be interpreted as "those queues are gone":
during a nameserver
* outage the topic route is unchanged, and revoking the whole assignment
would cause a mass
- * rebalance plus an offset rewind once the nameserver is back.
+ * rebalance plus an offset rewind once the nameserver is back. This covers
both the whole-round
+ * guards and the per-topic exclusion on partial failures.
*/
public class RocketMQSourceEnumeratorDiscoveryToleranceTest {
@@ -62,7 +63,7 @@ public class RocketMQSourceEnumeratorDiscoveryToleranceTest {
}
@Test
- public void partialFetchFailureKeepsFailedTopicQueuesTest() throws
Exception {
+ public void partialFetchFailureReportsFailedTopicsTest() throws Exception {
InnerConsumer consumer = mock(InnerConsumer.class);
when(consumer.fetchMessageQueues("t1"))
.thenReturn(
@@ -72,9 +73,11 @@ public class RocketMQSourceEnumeratorDiscoveryToleranceTest {
.thenReturn(failedFuture(new RuntimeException("nameserver
unreachable")));
RocketMQSourceEnumerator enumerator = newEnumerator(consumer, "t1;t2");
- Set<MessageQueue> discovered =
invokeRequestServiceDiscovery(enumerator);
+ RocketMQSourceEnumerator.ServiceDiscoveryResult result =
+ invokeRequestServiceDiscovery(enumerator);
- assertThat(discovered).containsExactlyInAnyOrder(queue("t1", 0),
queue("t1", 1));
+ assertThat(result.queues).containsExactlyInAnyOrder(queue("t1", 0),
queue("t1", 1));
+ assertThat(result.failedTopics).containsExactly("t2");
}
@Test
@@ -88,10 +91,51 @@ public class RocketMQSourceEnumeratorDiscoveryToleranceTest
{
.thenReturn(CompletableFuture.completedFuture(Arrays.asList(queue("t2", 0))));
RocketMQSourceEnumerator enumerator = newEnumerator(consumer, "t1;t2");
- Set<MessageQueue> discovered =
invokeRequestServiceDiscovery(enumerator);
+ RocketMQSourceEnumerator.ServiceDiscoveryResult result =
+ invokeRequestServiceDiscovery(enumerator);
- assertThat(discovered)
+ assertThat(result.queues)
.containsExactlyInAnyOrder(queue("t1", 0), queue("t1", 1),
queue("t2", 0));
+ assertThat(result.failedTopics).isEmpty();
+ }
+
+ @Test
+ public void partialFetchFailureDoesNotRevokeFailedTopicQueuesTest() throws
Exception {
+ InnerConsumer consumer = mock(InnerConsumer.class);
+ when(consumer.fetchMessageQueues("t1"))
+
.thenReturn(CompletableFuture.completedFuture(Arrays.asList(queue("t1", 0))));
+ when(consumer.fetchMessageQueues("t2"))
+ .thenReturn(failedFuture(new RuntimeException("nameserver
unreachable")));
+ RocketMQSourceEnumerator enumerator = newEnumerator(consumer, "t1;t2");
+ setAllocatedSet(enumerator, new HashSet<>(Arrays.asList(queue("t1",
0), queue("t2", 0))));
+
+ // The failed topic's queue is absent from the discovered set, but the
change result
+ // must be empty after exclusion, so the handler returns without
touching the
+ // assignment (and without scheduling any re-initialization on the
null context).
+ invokeHandleSourceQueueChange(enumerator,
invokeRequestServiceDiscovery(enumerator));
+
+ assertThat(allocatedSetOf(enumerator))
+ .containsExactlyInAnyOrder(queue("t1", 0), queue("t2", 0));
+ }
+
+ @Test
+ public void failedTopicExcludedFromDecreaseButGenuineDeletionKeptTest()
throws Exception {
+ InnerConsumer consumer = mock(InnerConsumer.class);
+ RocketMQSourceEnumerator enumerator = newEnumerator(consumer, "t1;t2");
+ setAllocatedSet(
+ enumerator,
+ new HashSet<>(Arrays.asList(queue("t1", 0), queue("t1", 1),
queue("t2", 0))));
+
+ // t1-q0 is really deleted (t1 fetched successfully without it); t2's
fetch failed,
+ // so t2-q0 must not be reported as decreased.
+ Object changeResult =
+ invokeGetSourceChangeResult(
+ enumerator,
+ new HashSet<>(Arrays.asList(queue("t1", 1))),
+ new HashSet<>(Arrays.asList("t2")));
+
+ assertThat(decreaseSetOf(changeResult)).containsExactly(queue("t1",
0));
+ assertThat(increaseSetOf(changeResult)).isEmpty();
}
@Test
@@ -100,12 +144,11 @@ public class
RocketMQSourceEnumeratorDiscoveryToleranceTest {
RocketMQSourceEnumerator enumerator = newEnumerator(consumer, "t1");
setAllocatedSet(enumerator, new HashSet<>(Arrays.asList(queue("t1",
0), queue("t1", 1))));
- Method handle =
- RocketMQSourceEnumerator.class.getDeclaredMethod(
- "handleSourceQueueChange", Set.class, Throwable.class);
- handle.setAccessible(true);
// Must return without touching the split assignment or scheduling any
re-initialization.
- handle.invoke(enumerator, new HashSet<MessageQueue>(), null);
+ invokeHandleSourceQueueChange(
+ enumerator,
+ new RocketMQSourceEnumerator.ServiceDiscoveryResult(
+ new HashSet<>(), new HashSet<>()));
assertThat(allocatedSetOf(enumerator)).hasSize(2);
}
@@ -118,15 +161,10 @@ public class
RocketMQSourceEnumeratorDiscoveryToleranceTest {
RocketMQSourceEnumerator enumerator = newEnumerator(consumer, "t1");
setAllocatedSet(enumerator, new HashSet<>(Arrays.asList(queue("t1",
0), queue("t1", 1))));
- Method diff =
- RocketMQSourceEnumerator.class.getDeclaredMethod(
- "getSourceChangeResult", Set.class);
- diff.setAccessible(true);
- Object changeResult = diff.invoke(enumerator, new
HashSet<MessageQueue>());
+ Object changeResult =
+ invokeGetSourceChangeResult(enumerator, new HashSet<>(), new
HashSet<>());
- Field decreaseField =
changeResult.getClass().getDeclaredField("decreaseSet");
- decreaseField.setAccessible(true);
- assertThat((Collection<?>) decreaseField.get(changeResult)).hasSize(2);
+ assertThat(decreaseSetOf(changeResult)).hasSize(2);
}
private static RocketMQSourceEnumerator newEnumerator(InnerConsumer
consumer, String topics)
@@ -146,12 +184,33 @@ public class
RocketMQSourceEnumeratorDiscoveryToleranceTest {
return enumerator;
}
- @SuppressWarnings("unchecked")
- private static Set<MessageQueue>
invokeRequestServiceDiscovery(RocketMQSourceEnumerator e)
- throws Exception {
+ private static RocketMQSourceEnumerator.ServiceDiscoveryResult
invokeRequestServiceDiscovery(
+ RocketMQSourceEnumerator e) throws Exception {
Method method =
RocketMQSourceEnumerator.class.getDeclaredMethod("requestServiceDiscovery");
method.setAccessible(true);
- return (Set<MessageQueue>) method.invoke(e);
+ return (RocketMQSourceEnumerator.ServiceDiscoveryResult)
method.invoke(e);
+ }
+
+ private static void invokeHandleSourceQueueChange(
+ RocketMQSourceEnumerator e,
RocketMQSourceEnumerator.ServiceDiscoveryResult result)
+ throws Exception {
+ Method handle =
+ RocketMQSourceEnumerator.class.getDeclaredMethod(
+ "handleSourceQueueChange",
+ RocketMQSourceEnumerator.ServiceDiscoveryResult.class,
+ Throwable.class);
+ handle.setAccessible(true);
+ handle.invoke(e, result, null);
+ }
+
+ private static Object invokeGetSourceChangeResult(
+ RocketMQSourceEnumerator e, Set<MessageQueue> latestSet,
Set<String> failedTopics)
+ throws Exception {
+ Method diff =
+ RocketMQSourceEnumerator.class.getDeclaredMethod(
+ "getSourceChangeResult", Set.class, Set.class);
+ diff.setAccessible(true);
+ return diff.invoke(e, latestSet, failedTopics);
}
@SuppressWarnings("unchecked")
@@ -170,6 +229,20 @@ public class
RocketMQSourceEnumeratorDiscoveryToleranceTest {
return ((java.util.Map<MessageQueue, Byte>) field.get(e)).keySet();
}
+ @SuppressWarnings("unchecked")
+ private static Collection<MessageQueue> decreaseSetOf(Object changeResult)
throws Exception {
+ Field field = changeResult.getClass().getDeclaredField("decreaseSet");
+ field.setAccessible(true);
+ return (Collection<MessageQueue>) field.get(changeResult);
+ }
+
+ @SuppressWarnings("unchecked")
+ private static Collection<MessageQueue> increaseSetOf(Object changeResult)
throws Exception {
+ Field field = changeResult.getClass().getDeclaredField("increaseSet");
+ field.setAccessible(true);
+ return (Collection<MessageQueue>) field.get(changeResult);
+ }
+
private static MessageQueue queue(String topic, int queueId) {
return new MessageQueue(topic, BROKER, queueId);
}
diff --git
a/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/source/enumerator/RocketMQSourceEnumeratorRevokedOffsetResumeTest.java
b/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/source/enumerator/RocketMQSourceEnumeratorRevokedOffsetResumeTest.java
new file mode 100644
index 0000000..abf1e44
--- /dev/null
+++
b/flink-connector-rocketmq/src/test/java/org/apache/flink/connector/rocketmq/source/enumerator/RocketMQSourceEnumeratorRevokedOffsetResumeTest.java
@@ -0,0 +1,241 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.connector.rocketmq.source.enumerator;
+
+import org.apache.flink.api.connector.source.Boundedness;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.connector.rocketmq.common.lock.SpinLock;
+import org.apache.flink.connector.rocketmq.source.InnerConsumer;
+import org.apache.flink.connector.rocketmq.source.RocketMQSourceOptions;
+import
org.apache.flink.connector.rocketmq.source.enumerator.offset.OffsetsSelector;
+import org.apache.flink.connector.rocketmq.source.split.RocketMQSourceSplit;
+
+import org.apache.rocketmq.common.message.MessageQueue;
+import org.junit.jupiter.api.Test;
+
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Field;
+import java.lang.reflect.Method;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/**
+ * A queue that was revoked and later reappears (broker rolling restart,
scale-down then scale-up)
+ * must resume from its last checkpointed offset, the same rule the migration
path applies to
+ * surviving queues. Re-resolving with the configured starting selector (here
EARLIEST) would replay
+ * the whole queue; only genuinely new queues may go through the selector.
+ */
+public class RocketMQSourceEnumeratorRevokedOffsetResumeTest {
+
+ private static final String BROKER = "broker-a";
+ private static final String TOPIC = "t1";
+
+ @Test
+ public void revokedQueueReaddedResumesFromCheckpointOffsetTest() throws
Exception {
+ InnerConsumer consumer = consumerWithOffsets(0L, 100L);
+ RocketMQSourceEnumerator enumerator = newEnumerator(consumer);
+ MessageQueue q0 = queue(0);
+ setAllocatedSet(enumerator, new HashSet<>(Arrays.asList(q0)));
+ checkedOffsetsOf(enumerator).put(q0, 42L);
+
+ Object decreaseResult =
+ invokeInitializeSourceSplits(
+ enumerator, new HashSet<>(), new
HashSet<>(Arrays.asList(q0)));
+ Set<RocketMQSourceSplit> decreaseSplits = splitsOf(decreaseResult,
"decreaseSet");
+ assertThat(decreaseSplits).hasSize(1);
+ assertThat(decreaseSplits.iterator().next().getIsIncrease()).isFalse();
+ assertThat(allocatedSetOf(enumerator)).isEmpty();
+ assertThat(revokedQueueOffsetsOf(enumerator)).containsEntry(q0, 42L);
+
+ Object increaseResult =
+ invokeInitializeSourceSplits(
+ enumerator, new HashSet<>(Arrays.asList(q0)), new
HashSet<>());
+ RocketMQSourceSplit readded = splitFor(splitsOf(increaseResult,
"increaseSet"), q0);
+ assertThat(readded.getIsIncrease()).isTrue();
+ assertThat(readded.getStartingOffset()).isEqualTo(42L);
+ assertThat(checkedOffsetsOf(enumerator)).containsEntry(q0, 42L);
+ assertThat(revokedQueueOffsetsOf(enumerator)).isEmpty();
+ }
+
+ @Test
+ public void genuinelyNewQueueUsesConfiguredSelectorTest() throws Exception
{
+ InnerConsumer consumer = consumerWithOffsets(0L, 100L);
+ RocketMQSourceEnumerator enumerator = newEnumerator(consumer);
+ MessageQueue q1 = queue(1);
+
+ Object increaseResult =
+ invokeInitializeSourceSplits(
+ enumerator, new HashSet<>(Arrays.asList(q1)), new
HashSet<>());
+
+ // Never allocated before: the EARLIEST selector resolves the min
offset (0).
+ assertThat(splitFor(splitsOf(increaseResult, "increaseSet"),
q1).getStartingOffset())
+ .isEqualTo(0L);
+ }
+
+ @Test
+ public void mixedReaddedAndNewQueuesResolveIndependentlyTest() throws
Exception {
+ InnerConsumer consumer = consumerWithOffsets(0L, 100L);
+ RocketMQSourceEnumerator enumerator = newEnumerator(consumer);
+ MessageQueue q0 = queue(0);
+ MessageQueue q1 = queue(1);
+ setAllocatedSet(enumerator, new HashSet<>(Arrays.asList(q0)));
+ checkedOffsetsOf(enumerator).put(q0, 42L);
+ invokeInitializeSourceSplits(enumerator, new HashSet<>(), new
HashSet<>(Arrays.asList(q0)));
+
+ Object increaseResult =
+ invokeInitializeSourceSplits(
+ enumerator, new HashSet<>(Arrays.asList(q0, q1)), new
HashSet<>());
+
+ Set<RocketMQSourceSplit> increaseSplits = splitsOf(increaseResult,
"increaseSet");
+ assertThat(splitFor(increaseSplits,
q0).getStartingOffset()).isEqualTo(42L);
+ assertThat(splitFor(increaseSplits,
q1).getStartingOffset()).isEqualTo(0L);
+ }
+
+ @Test
+ public void nonCheckpointedOffsetFallsBackToSelectorTest() throws
Exception {
+ InnerConsumer consumer = consumerWithOffsets(0L, 100L);
+ RocketMQSourceEnumerator enumerator = newEnumerator(consumer);
+ MessageQueue q0 = queue(0);
+ setAllocatedSet(enumerator, new HashSet<>(Arrays.asList(q0)));
+ // NO_STOPPING_OFFSET: e.g. restored from a checkpoint before the
first offset report.
+ // There is no meaningful checkpointed progress to retain.
+ checkedOffsetsOf(enumerator).put(q0,
RocketMQSourceSplit.NO_STOPPING_OFFSET);
+
+ invokeInitializeSourceSplits(enumerator, new HashSet<>(), new
HashSet<>(Arrays.asList(q0)));
+ assertThat(revokedQueueOffsetsOf(enumerator)).isEmpty();
+
+ Object increaseResult =
+ invokeInitializeSourceSplits(
+ enumerator, new HashSet<>(Arrays.asList(q0)), new
HashSet<>());
+ assertThat(splitFor(splitsOf(increaseResult, "increaseSet"),
q0).getStartingOffset())
+ .isEqualTo(0L);
+ }
+
+ private static RocketMQSourceEnumerator newEnumerator(InnerConsumer
consumer) throws Exception {
+ Configuration configuration = new Configuration();
+ configuration.set(RocketMQSourceOptions.TOPIC, TOPIC);
+ RocketMQSourceEnumerator enumerator =
+ new RocketMQSourceEnumerator(
+ OffsetsSelector.earliest(),
+ OffsetsSelector.latest(),
+ Boundedness.CONTINUOUS_UNBOUNDED,
+ configuration,
+ null);
+ Field consumerField =
RocketMQSourceEnumerator.class.getDeclaredField("consumer");
+ consumerField.setAccessible(true);
+ consumerField.set(enumerator, consumer);
+ return enumerator;
+ }
+
+ private static InnerConsumer consumerWithOffsets(long minOffset, long
maxOffset) {
+ InnerConsumer consumer = mock(InnerConsumer.class);
+ when(consumer.seekMinOffset(any(MessageQueue.class)))
+ .thenReturn(CompletableFuture.completedFuture(minOffset));
+ when(consumer.seekMaxOffset(any(MessageQueue.class)))
+ .thenReturn(CompletableFuture.completedFuture(maxOffset));
+ return consumer;
+ }
+
+ /**
+ * Invokes the private initialization step directly. The production path
releases the enumerator
+ * lock in {@code handleSplitChanges}; bypassing that handler means the
test must release it
+ * here, otherwise the next round spins forever on the non-reentrant lock.
+ */
+ private static Object invokeInitializeSourceSplits(
+ RocketMQSourceEnumerator e, Set<MessageQueue> increase,
Set<MessageQueue> decrease)
+ throws Exception {
+ Class<?> changeClass =
+ Class.forName(
+
"org.apache.flink.connector.rocketmq.source.enumerator.RocketMQSourceEnumerator$SourceChangeResult");
+ Constructor<?> constructor =
changeClass.getDeclaredConstructor(Set.class, Set.class);
+ constructor.setAccessible(true);
+ Object changeResult = constructor.newInstance(increase, decrease);
+ Method method =
+ RocketMQSourceEnumerator.class.getDeclaredMethod(
+ "initializeSourceSplits", changeClass);
+ method.setAccessible(true);
+ try {
+ return method.invoke(e, changeResult);
+ } finally {
+ Field lockField =
RocketMQSourceEnumerator.class.getDeclaredField("lock");
+ lockField.setAccessible(true);
+ ((SpinLock) lockField.get(e)).unlock();
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private static Set<RocketMQSourceSplit> splitsOf(Object splitChangeResult,
String fieldName)
+ throws Exception {
+ Field field = splitChangeResult.getClass().getDeclaredField(fieldName);
+ field.setAccessible(true);
+ return (Set<RocketMQSourceSplit>) field.get(splitChangeResult);
+ }
+
+ private static RocketMQSourceSplit splitFor(
+ Set<RocketMQSourceSplit> splits, MessageQueue messageQueue) {
+ return splits.stream()
+ .filter(split -> split.getMessageQueue().equals(messageQueue))
+ .findFirst()
+ .orElseThrow(() -> new AssertionError("No split for " +
messageQueue));
+ }
+
+ @SuppressWarnings("unchecked")
+ private static void setAllocatedSet(RocketMQSourceEnumerator e,
Set<MessageQueue> queues)
+ throws Exception {
+ Field field =
RocketMQSourceEnumerator.class.getDeclaredField("allocatedSet");
+ field.setAccessible(true);
+ Map<MessageQueue, Byte> allocatedSet = (Map<MessageQueue, Byte>)
field.get(e);
+ queues.forEach(q -> allocatedSet.put(q, (byte) 0));
+ }
+
+ @SuppressWarnings("unchecked")
+ private static Set<MessageQueue> allocatedSetOf(RocketMQSourceEnumerator
e) throws Exception {
+ Field field =
RocketMQSourceEnumerator.class.getDeclaredField("allocatedSet");
+ field.setAccessible(true);
+ return ((Map<MessageQueue, Byte>) field.get(e)).keySet();
+ }
+
+ @SuppressWarnings("unchecked")
+ private static Map<MessageQueue, Long>
checkedOffsetsOf(RocketMQSourceEnumerator e)
+ throws Exception {
+ Field field =
RocketMQSourceEnumerator.class.getDeclaredField("checkedOffsets");
+ field.setAccessible(true);
+ return (Map<MessageQueue, Long>) field.get(e);
+ }
+
+ @SuppressWarnings("unchecked")
+ private static Map<MessageQueue, Long>
revokedQueueOffsetsOf(RocketMQSourceEnumerator e)
+ throws Exception {
+ Field field =
RocketMQSourceEnumerator.class.getDeclaredField("revokedQueueOffsets");
+ field.setAccessible(true);
+ return (Map<MessageQueue, Long>) field.get(e);
+ }
+
+ private static MessageQueue queue(int queueId) {
+ return new MessageQueue(TOPIC, BROKER, queueId);
+ }
+}
