This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new f8bfcffc9ba Fix subscription stalls after consumer disconnect (#18760)
f8bfcffc9ba is described below
commit f8bfcffc9baf0e1b3c046bf8bf9a2a8b4b01a74b
Author: Caideyipi <[email protected]>
AuthorDate: Wed Sep 30 10:08:35 2026 +0800
Fix subscription stalls after consumer disconnect (#18760)
---
.../agent/SubscriptionBrokerAgent.java | 8 ++++
.../agent/SubscriptionReceiverAgent.java | 12 +++++
.../broker/ConsensusSubscriptionBroker.java | 8 ++++
.../subscription/broker/ISubscriptionBroker.java | 3 ++
.../db/subscription/broker/SubscriptionBroker.java | 7 +++
.../broker/SubscriptionPrefetchingQueue.java | 11 +++++
.../consensus/ConsensusPrefetchingQueue.java | 10 +++++
.../receiver/SubscriptionReceiver.java | 5 +++
.../receiver/SubscriptionReceiverV1.java | 15 +++++++
.../agent/SubscriptionReceiverAgentTest.java | 51 ++++++++++++++++++++++
...onsensusSubscriptionBrokerPayloadLimitTest.java | 14 ++++++
.../SubscriptionBrokerAgentPayloadLimitTest.java | 14 ++++++
.../receiver/SubscriptionReceiverV1Test.java | 22 +++++++++-
13 files changed, 178 insertions(+), 2 deletions(-)
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionBrokerAgent.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionBrokerAgent.java
index e65f1d88041..93f64aa0381 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionBrokerAgent.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionBrokerAgent.java
@@ -228,6 +228,14 @@ public class SubscriptionBrokerAgent {
return false;
}
+ public int requeueInFlightEvents(final ConsumerConfig consumerConfig) {
+ final String consumerGroupId = consumerConfig.getConsumerGroupId();
+ final String consumerId = consumerConfig.getConsumerId();
+ return getBrokers(consumerGroupId).stream()
+ .mapToInt(broker -> broker.requeueInFlightEvents(consumerId))
+ .sum();
+ }
+
public int refreshInFlightEventLeases(
final ConsumerConfig consumerConfig,
final List<SubscriptionCommitContext> processorBufferedCommitContexts) {
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionReceiverAgent.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionReceiverAgent.java
index dc42d4fbcae..a213eac6dfc 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionReceiverAgent.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionReceiverAgent.java
@@ -161,6 +161,8 @@ public class SubscriptionReceiverAgent {
if (isHandshake(req)) {
if (isSuccessful(requestResult.response)) {
if (currentReceiver != null && currentReceiver != receiver) {
+ releaseInFlightEventsOnNewInstanceTakeover(
+ currentReceiver, receiver.getConsumerInstanceId());
invalidateReplacedReceiver(currentReceiver, identity);
}
return receiver;
@@ -319,6 +321,8 @@ public class SubscriptionReceiverAgent {
}
if (receiver.getConsumerInstanceId() != null
&& !shouldKeepCurrentReceiver(currentReceiver,
receiver.getConsumerInstanceId())) {
+ releaseInFlightEventsOnNewInstanceTakeover(
+ currentReceiver, receiver.getConsumerInstanceId());
invalidateReplacedReceiver(currentReceiver, key);
registered.set(true);
return receiver;
@@ -361,6 +365,14 @@ public class SubscriptionReceiverAgent {
receiver.invalidateConsumer();
}
+ private static void releaseInFlightEventsOnNewInstanceTakeover(
+ final SubscriptionReceiver currentReceiver, final String
incomingConsumerInstanceId) {
+ if (incomingConsumerInstanceId != null
+ && !Objects.equals(currentReceiver.getConsumerInstanceId(),
incomingConsumerInstanceId)) {
+ currentReceiver.releaseInFlightEvents();
+ }
+ }
+
private static ConsumerConnection getConsumerConnection(
final TPipeSubscribeReq req, final SubscriptionReceiver receiver) {
if (isHandshake(req) && req.isSetBody()) {
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBroker.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBroker.java
index 9765cdc4ea0..6b29f069524 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBroker.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBroker.java
@@ -316,6 +316,14 @@ public class ConsensusSubscriptionBroker implements
ISubscriptionBroker {
return Objects.nonNull(queue) && queue.requeue(consumerId, commitContext);
}
+ @Override
+ public int requeueInFlightEvents(final String consumerId) {
+ return topicNameToConsensusPrefetchingQueues.values().stream()
+ .flatMap(List::stream)
+ .mapToInt(queue -> queue.requeueInFlightEvents(consumerId))
+ .sum();
+ }
+
@Override
public int refreshInFlightEventLeases(
final String consumerId, final List<SubscriptionCommitContext>
commitContexts) {
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ISubscriptionBroker.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ISubscriptionBroker.java
index 90eb947292e..567b93efada 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ISubscriptionBroker.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ISubscriptionBroker.java
@@ -52,6 +52,9 @@ public interface ISubscriptionBroker {
/** Returns an in-flight event to its prefetching queue without incrementing
its nack count. */
boolean requeue(String consumerId, SubscriptionCommitContext commitContext);
+ /** Returns all in-flight events held by a consumer without incrementing
their nack counts. */
+ int requeueInFlightEvents(String consumerId);
+
default List<SubscriptionCommitContext> selectAcceptedCommitContexts(
final List<SubscriptionCommitContext> commitContexts) {
if (Objects.isNull(commitContexts) || commitContexts.isEmpty()) {
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionBroker.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionBroker.java
index 0353e3d74b1..3b1dec89dd1 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionBroker.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionBroker.java
@@ -392,6 +392,13 @@ public class SubscriptionBroker implements
ISubscriptionBroker {
&& prefetchingQueue.requeue(consumerId, commitContext);
}
+ @Override
+ public int requeueInFlightEvents(final String consumerId) {
+ return topicNameToPrefetchingQueue.values().stream()
+ .mapToInt(queue -> queue.requeueInFlightEvents(consumerId))
+ .sum();
+ }
+
@Override
public int refreshInFlightEventLeases(
final String consumerId, final List<SubscriptionCommitContext>
commitContexts) {
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java
index 2f37ac6f64b..b760ee62f6f 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java
@@ -850,6 +850,17 @@ public abstract class SubscriptionPrefetchingQueue {
}
}
+ public int requeueInFlightEvents(final String consumerId) {
+ int requeuedCount = 0;
+ for (final Pair<String, SubscriptionCommitContext> key :
+ ImmutableSet.copyOf(inFlightEvents.keySet())) {
+ if (Objects.equals(consumerId, key.left) && requeue(consumerId,
key.right)) {
+ requeuedCount++;
+ }
+ }
+ return requeuedCount;
+ }
+
/**
* @return {@code true} if ack successfully
*/
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java
index 79cb048afac..d415656986b 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java
@@ -2717,6 +2717,16 @@ public class ConsensusPrefetchingQueue {
}
}
+ public int requeueInFlightEvents(final String consumerId) {
+ int requeuedCount = 0;
+ for (final InFlightEventKey key : new
ArrayList<>(inFlightEvents.keySet())) {
+ if (Objects.equals(consumerId, key.consumerId) && requeue(consumerId,
key.commitContext)) {
+ requeuedCount++;
+ }
+ }
+ return requeuedCount;
+ }
+
private boolean canAcceptCommitContext(
final SubscriptionCommitContext commitContext, final String action,
final boolean silent) {
if (isClosed || closeRequested || pendingSeekRequest != null) {
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiver.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiver.java
index 1d2fbe4fc18..01c736b29dc 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiver.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiver.java
@@ -33,6 +33,11 @@ public interface SubscriptionReceiver {
void handleExit();
+ /** Releases the in-flight events held by this receiver's current consumer.
*/
+ default void releaseInFlightEvents() {
+ // no-op for receivers without subscription event ownership
+ }
+
void handleTimeout();
/**
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java
index 57c768e4c7c..2c2b00eca12 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java
@@ -232,6 +232,16 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
}
}
+ @Override
+ public void releaseInFlightEvents() {
+ synchronized (this) {
+ final ConsumerConfig consumerConfig = sharedConsumerConfig;
+ if (Objects.nonNull(consumerConfig) && !consumerInvalidated) {
+ releaseInFlightEvents(consumerConfig);
+ }
+ }
+ }
+
@Override
public void handleExit() {
synchronized (this) {
@@ -246,6 +256,7 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
// progress to continue consuming. When another connection has
already taken ownership,
// the receiver is invalidated and even this best-effort cleanup
must be skipped.
if (!consumerInvalidated) {
+ releaseInFlightEvents(consumerConfig);
// When handling exit, unsubscribe from topics that have already
been completed as much
// as possible to release some resources (such as the underlying
pipe) in a timely
// manner.
@@ -493,6 +504,10 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
.refreshInFlightEventLeases(consumerConfig,
processorBufferedCommitContexts);
}
+ protected int releaseInFlightEvents(final ConsumerConfig consumerConfig) {
+ return SubscriptionAgent.broker().requeueInFlightEvents(consumerConfig);
+ }
+
private TPipeSubscribeResp handlePipeSubscribeSubscribe(final
PipeSubscribeSubscribeReq req) {
try {
return handlePipeSubscribeSubscribeInternal(req);
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/agent/SubscriptionReceiverAgentTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/agent/SubscriptionReceiverAgentTest.java
index 2d48518e5f2..21ef6ecb101 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/agent/SubscriptionReceiverAgentTest.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/agent/SubscriptionReceiverAgentTest.java
@@ -331,6 +331,51 @@ public class SubscriptionReceiverAgentTest {
Assert.assertFalse(receivers.get(1).invalidated);
}
+ @Test
+ public void testNewConsumerInstanceTakeoverReleasesOldInFlightEvents()
throws Exception {
+ final CopyOnWriteArrayList<FakeSubscriptionReceiver> receivers = new
CopyOnWriteArrayList<>();
+ final SubscriptionReceiverAgent agent = createAgent(receivers, false /*
closeOnTimeout */);
+
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ agent
+ .handle(createHandshakeRequest("group", "consumer", "instance-1"),
"root")
+ .getStatus()
+ .getCode());
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ handleOnNewConnection(agent, createHandshakeRequest("group",
"consumer", "instance-2"))
+ .getStatus()
+ .getCode());
+
+ Assert.assertEquals(2, receivers.size());
+ Assert.assertEquals(1, receivers.get(0).releaseCount.get());
+ Assert.assertTrue(receivers.get(0).invalidated);
+ Assert.assertEquals(0, receivers.get(1).releaseCount.get());
+ }
+
+ @Test
+ public void testSameConsumerInstanceReconnectPreservesInFlightEvents()
throws Exception {
+ final CopyOnWriteArrayList<FakeSubscriptionReceiver> receivers = new
CopyOnWriteArrayList<>();
+ final SubscriptionReceiverAgent agent = createAgent(receivers, false /*
closeOnTimeout */);
+
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ agent
+ .handle(createHandshakeRequest("group", "consumer", "instance-1"),
"root")
+ .getStatus()
+ .getCode());
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ handleOnNewConnection(agent, createHandshakeRequest("group",
"consumer", "instance-1"))
+ .getStatus()
+ .getCode());
+
+ Assert.assertEquals(2, receivers.size());
+ Assert.assertEquals(0, receivers.get(0).releaseCount.get());
+ Assert.assertTrue(receivers.get(0).invalidated);
+ }
+
@Test
public void testLateHandshakeCannotTakeOverNewReceiver() throws Exception {
final CopyOnWriteArrayList<FakeSubscriptionReceiver> receivers = new
CopyOnWriteArrayList<>();
@@ -565,6 +610,7 @@ public class SubscriptionReceiverAgentTest {
private final CountDownLatch releaseHandshake;
private final AtomicInteger timeoutCount = new AtomicInteger();
private final AtomicInteger exitCount = new AtomicInteger();
+ private final AtomicInteger releaseCount = new AtomicInteger();
private volatile ConsumerConfig consumerConfig;
private volatile boolean invalidated;
@@ -627,6 +673,11 @@ public class SubscriptionReceiverAgentTest {
exitCount.incrementAndGet();
}
+ @Override
+ public void releaseInFlightEvents() {
+ releaseCount.incrementAndGet();
+ }
+
@Override
public void handleTimeout() {
timeoutCount.incrementAndGet();
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerPayloadLimitTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerPayloadLimitTest.java
index 1d2e7137535..d07ec0b557f 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerPayloadLimitTest.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerPayloadLimitTest.java
@@ -73,6 +73,20 @@ public class ConsensusSubscriptionBrokerPayloadLimitTest {
verify(secondQueue).requeue(CONSUMER_ID, secondCommitContext);
}
+ @Test
+ public void testRequeueInFlightEventsAcrossRegionQueues() throws Exception {
+ final ConsensusSubscriptionBroker broker = new
ConsensusSubscriptionBroker(CONSUMER_GROUP_ID);
+ final ConsensusPrefetchingQueue firstQueue =
mock(ConsensusPrefetchingQueue.class);
+ final ConsensusPrefetchingQueue secondQueue =
mock(ConsensusPrefetchingQueue.class);
+ when(firstQueue.requeueInFlightEvents(CONSUMER_ID)).thenReturn(2);
+ when(secondQueue.requeueInFlightEvents(CONSUMER_ID)).thenReturn(3);
+ bindQueues(broker, Arrays.asList(firstQueue, secondQueue));
+
+ assertEquals(5, broker.requeueInFlightEvents(CONSUMER_ID));
+ verify(firstQueue).requeueInFlightEvents(CONSUMER_ID);
+ verify(secondQueue).requeueInFlightEvents(CONSUMER_ID);
+ }
+
private static SubscriptionCommitContext newCommitContext(
final int regionId, final int commitId) {
return new SubscriptionCommitContext(
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionBrokerAgentPayloadLimitTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionBrokerAgentPayloadLimitTest.java
index 782f2338e9d..d0fdafc6e4d 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionBrokerAgentPayloadLimitTest.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionBrokerAgentPayloadLimitTest.java
@@ -78,6 +78,20 @@ public class SubscriptionBrokerAgentPayloadLimitTest {
assertEquals(0L, secondEvent.getNackCount());
}
+ @Test
+ public void testRequeueInFlightEventsAcrossBrokers() throws Exception {
+ final SubscriptionBrokerAgent agent = new SubscriptionBrokerAgent();
+ final ISubscriptionBroker firstBroker = mock(ISubscriptionBroker.class);
+ final ISubscriptionBroker secondBroker = mock(ISubscriptionBroker.class);
+ when(firstBroker.requeueInFlightEvents(CONSUMER_ID)).thenReturn(2);
+ when(secondBroker.requeueInFlightEvents(CONSUMER_ID)).thenReturn(3);
+ bindBrokers(agent, firstBroker, secondBroker);
+
+ assertEquals(5, agent.requeueInFlightEvents(createConsumerConfig()));
+ verify(firstBroker).requeueInFlightEvents(CONSUMER_ID);
+ verify(secondBroker).requeueInFlightEvents(CONSUMER_ID);
+ }
+
private static SubscriptionEvent newEvent(final int commitId) {
return new SubscriptionEvent(
SubscriptionPollResponseType.TERMINATION.getType(),
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1Test.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1Test.java
index 53f42666c65..8b2bd022241 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1Test.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1Test.java
@@ -131,7 +131,15 @@ public class SubscriptionReceiverV1Test {
@Test
@SuppressWarnings("unchecked")
public void testHandleExitKeepsSharedConsumerStateForTimeoutCleanup() throws
Exception {
- final SubscriptionReceiverV1 receiver = new SubscriptionReceiverV1();
+ final AtomicLong releaseCount = new AtomicLong();
+ final SubscriptionReceiverV1 receiver =
+ new SubscriptionReceiverV1() {
+ @Override
+ protected int releaseInFlightEvents(final ConsumerConfig
consumerConfig) {
+ releaseCount.incrementAndGet();
+ return 1;
+ }
+ };
final ConsumerConfig consumerConfig = createConsumerConfig(1_000L);
setField(receiver, "sharedConsumerConfig", consumerConfig);
final ThreadLocal<ConsumerConfig> consumerConfigThreadLocal =
@@ -144,12 +152,21 @@ public class SubscriptionReceiverV1Test {
Assert.assertFalse((boolean) getField(receiver, "consumerInvalidated"));
Assert.assertTrue(receiver.hasActiveConsumer());
Assert.assertNull(consumerConfigThreadLocal.get());
+ Assert.assertEquals(1L, releaseCount.get());
}
@Test
@SuppressWarnings("unchecked")
public void testHandleExitClearsThreadLocalStateAfterInvalidation() throws
Exception {
- final SubscriptionReceiverV1 receiver = new SubscriptionReceiverV1();
+ final AtomicLong releaseCount = new AtomicLong();
+ final SubscriptionReceiverV1 receiver =
+ new SubscriptionReceiverV1() {
+ @Override
+ protected int releaseInFlightEvents(final ConsumerConfig
consumerConfig) {
+ releaseCount.incrementAndGet();
+ return 1;
+ }
+ };
final ConsumerConfig consumerConfig = createConsumerConfig(1_000L);
setField(receiver, "sharedConsumerConfig", consumerConfig);
final ThreadLocal<ConsumerConfig> consumerConfigThreadLocal =
@@ -162,6 +179,7 @@ public class SubscriptionReceiverV1Test {
Assert.assertFalse(receiver.hasActiveConsumer());
Assert.assertTrue((boolean) getField(receiver, "consumerInvalidated"));
Assert.assertNull(consumerConfigThreadLocal.get());
+ Assert.assertEquals(0L, releaseCount.get());
}
@Test
