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 2479dc63531 Subscription: clean up consumers after abrupt disconnect
(#18434)
2479dc63531 is described below
commit 2479dc63531a38db95b8249b1b4589cc98e056d9
Author: Caideyipi <[email protected]>
AuthorDate: Wed Aug 12 10:23:11 2026 +0800
Subscription: clean up consumers after abrupt disconnect (#18434)
* Subscription: clean up disconnected consumers after timeout
* Test: cover hard-disconnected subscription consumer cleanup
* Test: verify timeout cleanup after in-flight request
---
...ensusSubscriptionConsumerDisconnectTableIT.java | 298 +++++++++++++++++++++
.../agent/SubscriptionReceiverAgent.java | 242 +++++++++++++++--
.../receiver/SubscriptionReceiver.java | 21 ++
.../receiver/SubscriptionReceiverV1.java | 115 +++++---
.../agent/SubscriptionReceiverAgentTest.java | 273 +++++++++++++++++++
.../receiver/SubscriptionReceiverV1Test.java | 94 ++++++-
6 files changed, 974 insertions(+), 69 deletions(-)
diff --git
a/integration-test/src/test/java/org/apache/iotdb/subscription/it/consensus/local/tablemodel/IoTDBConsensusSubscriptionConsumerDisconnectTableIT.java
b/integration-test/src/test/java/org/apache/iotdb/subscription/it/consensus/local/tablemodel/IoTDBConsensusSubscriptionConsumerDisconnectTableIT.java
new file mode 100644
index 00000000000..67b30779ce9
--- /dev/null
+++
b/integration-test/src/test/java/org/apache/iotdb/subscription/it/consensus/local/tablemodel/IoTDBConsensusSubscriptionConsumerDisconnectTableIT.java
@@ -0,0 +1,298 @@
+/*
+ * 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.iotdb.subscription.it.consensus.local.tablemodel;
+
+import org.apache.iotdb.isession.ITableSession;
+import org.apache.iotdb.it.env.EnvFactory;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.itbase.category.LocalStandaloneIT;
+import org.apache.iotdb.itbase.category.ManualIT;
+import org.apache.iotdb.session.subscription.ISubscriptionTableSession;
+import org.apache.iotdb.session.subscription.SubscriptionSessionWrapper;
+import org.apache.iotdb.session.subscription.SubscriptionTableSessionBuilder;
+import
org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumer;
+import
org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumerBuilder;
+import org.apache.iotdb.session.subscription.model.Subscription;
+import
org.apache.iotdb.subscription.it.consensus.local.AbstractSubscriptionConsensusLocalIT;
+
+import org.apache.thrift.transport.TTransport;
+import org.awaitility.Awaitility;
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+
+import java.lang.reflect.Field;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+@RunWith(IoTDBTestRunner.class)
+@Category({LocalStandaloneIT.class, ManualIT.class})
+public class IoTDBConsensusSubscriptionConsumerDisconnectTableIT
+ extends AbstractSubscriptionConsensusLocalIT {
+
+ private static final long HEARTBEAT_INTERVAL_MS = 1_000L;
+ private static final int INITIAL_ROW_COUNT = 20;
+ private static final int ROW_COUNT_DURING_DISCONNECT = 200;
+
+ @Test
+ public void testOtherConsumersContinueAfterConsumerHardDisconnect() throws
Exception {
+ final ConsensusSubscriptionTableITSupport.TestIdentifiers keepOneIds =
+
ConsensusSubscriptionTableITSupport.newIdentifiers("hard_disconnect_keep_one");
+ final ConsensusSubscriptionTableITSupport.TestIdentifiers keepTwoIds =
+
ConsensusSubscriptionTableITSupport.newIdentifiers("hard_disconnect_keep_two");
+ final ConsensusSubscriptionTableITSupport.TestIdentifiers stoppedIds =
+
ConsensusSubscriptionTableITSupport.newIdentifiers("hard_disconnect_stopped");
+ final String database = keepOneIds.getDatabase();
+ final String table = "t1";
+ final Set<String> topics = new LinkedHashSet<>();
+ topics.add(keepOneIds.getTopic());
+ topics.add(keepTwoIds.getTopic());
+ topics.add(stoppedIds.getTopic());
+
+ SubscriptionTablePullConsumer keepOneConsumer = null;
+ SubscriptionTablePullConsumer keepTwoConsumer = null;
+ SubscriptionTablePullConsumer stoppedConsumer = null;
+ final ExecutorService writer = Executors.newSingleThreadExecutor();
+ final CountDownLatch disconnectPointReached = new CountDownLatch(1);
+ final CountDownLatch disconnectCompleted = new CountDownLatch(1);
+ Future<Set<String>> writeFuture = null;
+
+ try {
+ ConsensusSubscriptionTableITSupport.bootstrapDatabaseAndTable(
+ database, table,
ConsensusSubscriptionTableITSupport.DEFAULT_TABLE_SCHEMA);
+ for (final String topic : topics) {
+ ConsensusSubscriptionTableITSupport.createConsensusTopic(topic,
database, table);
+ }
+
+ keepOneConsumer = createConsumer(keepOneIds.getConsumerId(),
keepOneIds.getConsumerGroupId());
+ keepTwoConsumer = createConsumer(keepTwoIds.getConsumerId(),
keepTwoIds.getConsumerGroupId());
+ stoppedConsumer = createConsumer(stoppedIds.getConsumerId(),
stoppedIds.getConsumerGroupId());
+ keepOneConsumer.subscribe(keepOneIds.getTopic());
+ keepTwoConsumer.subscribe(keepTwoIds.getTopic());
+ stoppedConsumer.subscribe(stoppedIds.getTopic());
+
+ awaitSubscriptionPresent(
+ keepOneIds.getTopic(), keepOneIds.getConsumerGroupId(),
keepOneIds.getConsumerId());
+ awaitSubscriptionPresent(
+ keepTwoIds.getTopic(), keepTwoIds.getConsumerGroupId(),
keepTwoIds.getConsumerId());
+ awaitSubscriptionPresent(
+ stoppedIds.getTopic(), stoppedIds.getConsumerGroupId(),
stoppedIds.getConsumerId());
+
+ final Set<String> initialRows =
+ ConsensusSubscriptionTableITSupport.insertRows(
+ database, table, 100L, INITIAL_ROW_COUNT, true);
+ assertRowsConsumed(keepOneConsumer, initialRows);
+ assertRowsConsumed(keepTwoConsumer, initialRows);
+ assertRowsConsumed(stoppedConsumer, initialRows);
+
+ writeFuture =
+ writer.submit(
+ () ->
+ insertRowsAcrossDisconnect(
+ database, table, 1_000L, disconnectPointReached,
disconnectCompleted));
+ Assert.assertTrue(
+ "Timed out waiting for the batch writer to reach the disconnect
point",
+ disconnectPointReached.await(30, TimeUnit.SECONDS));
+
+ disconnectWithoutClose(stoppedConsumer);
+ disconnectCompleted.countDown();
+ final Set<String> rowsWrittenAcrossDisconnect = writeFuture.get(2,
TimeUnit.MINUTES);
+
+ awaitSubscriptionRemoved(stoppedIds.getTopic());
+ awaitSubscriptionPresent(
+ keepOneIds.getTopic(), keepOneIds.getConsumerGroupId(),
keepOneIds.getConsumerId());
+ awaitSubscriptionPresent(
+ keepTwoIds.getTopic(), keepTwoIds.getConsumerGroupId(),
keepTwoIds.getConsumerId());
+
+ assertRowsConsumed(keepOneConsumer, rowsWrittenAcrossDisconnect);
+ assertRowsConsumed(keepTwoConsumer, rowsWrittenAcrossDisconnect);
+ } finally {
+ disconnectCompleted.countDown();
+ if (writeFuture != null && !writeFuture.isDone()) {
+ writeFuture.cancel(true);
+ }
+ writer.shutdownNow();
+ writer.awaitTermination(10, TimeUnit.SECONDS);
+ closeQuietly(stoppedConsumer);
+ closeQuietly(keepTwoConsumer);
+ closeQuietly(keepOneConsumer);
+ ConsensusSubscriptionTableITSupport.cleanup(null, topics, database);
+ }
+ }
+
+ private static SubscriptionTablePullConsumer createConsumer(
+ final String consumerId, final String consumerGroupId) throws Exception {
+ final SubscriptionTablePullConsumer consumer =
+ (SubscriptionTablePullConsumer)
+ new SubscriptionTablePullConsumerBuilder()
+ .host(EnvFactory.getEnv().getIP())
+ .port(Integer.parseInt(EnvFactory.getEnv().getPort()))
+ .consumerId(consumerId)
+ .consumerGroupId(consumerGroupId)
+ .heartbeatIntervalMs(HEARTBEAT_INTERVAL_MS)
+ .endpointsSyncIntervalMs(5_000L)
+ .autoCommit(false)
+ .build();
+ consumer.open();
+ return consumer;
+ }
+
+ private static Set<String> insertRowsAcrossDisconnect(
+ final String database,
+ final String table,
+ final long startTimestamp,
+ final CountDownLatch disconnectPointReached,
+ final CountDownLatch disconnectCompleted)
+ throws Exception {
+ final Set<String> rowKeys = new LinkedHashSet<>();
+ try (final ITableSession session =
EnvFactory.getEnv().getTableSessionConnection()) {
+ session.executeNonQueryStatement("use " + database);
+ for (int row = 0; row < ROW_COUNT_DURING_DISCONNECT; row++) {
+ if (row == ROW_COUNT_DURING_DISCONNECT / 4) {
+ disconnectPointReached.countDown();
+ Assert.assertTrue(
+ "Timed out waiting for the consumer hard disconnect",
+ disconnectCompleted.await(30, TimeUnit.SECONDS));
+ }
+ final long timestamp = startTimestamp + row;
+ session.executeNonQueryStatement(
+ String.format(
+ Locale.ROOT,
+ "insert into %s(tag1, s1, time) values ('batch', %d, %d)",
+ table,
+ timestamp * 10L,
+ timestamp));
+ rowKeys.add(ConsensusSubscriptionTableITSupport.rowKey(database,
table, timestamp));
+ }
+ session.executeNonQueryStatement("flush");
+ }
+ return rowKeys;
+ }
+
+ private static void disconnectWithoutClose(final
SubscriptionTablePullConsumer consumer)
+ throws Exception {
+ final AtomicBoolean isClosed = (AtomicBoolean) getField(consumer,
"isClosed");
+ isClosed.set(true);
+
+ // Let the scheduled workers observe isClosed before severing the only
provider connection.
+ ConsensusSubscriptionTableITSupport.pause(HEARTBEAT_INTERVAL_MS * 2L);
+
+ final Object providers = getField(consumer, "providers");
+ final Map<?, ?> providerMap = (Map<?, ?>) getField(providers,
"subscriptionProviders");
+ final List<?> providerSnapshot = new ArrayList<>(providerMap.values());
+ Assert.assertFalse("Expected at least one subscription provider",
providerSnapshot.isEmpty());
+ for (final Object provider : providerSnapshot) {
+ final SubscriptionSessionWrapper session =
+ (SubscriptionSessionWrapper) getField(provider, "session");
+ final Object connection = session.getSessionConnection();
+ final TTransport transport = (TTransport) getField(connection,
"transport");
+ Assert.assertTrue("Expected an open subscription transport",
transport.isOpen());
+ transport.close();
+ }
+ }
+
+ private static void awaitSubscriptionPresent(
+ final String topicName, final String consumerGroupId, final String
consumerId)
+ throws Exception {
+ try (final ISubscriptionTableSession session =
createSubscriptionSession()) {
+ Awaitility.await()
+ .pollInSameThread()
+ .pollInterval(Duration.ofMillis(500))
+ .atMost(Duration.ofSeconds(30))
+ .untilAsserted(
+ () -> {
+ final Set<Subscription> subscriptions =
session.getSubscriptions(topicName);
+ Assert.assertEquals(subscriptions.toString(), 1,
subscriptions.size());
+ final Subscription subscription =
subscriptions.iterator().next();
+ Assert.assertEquals(consumerGroupId,
subscription.getConsumerGroupId());
+ Assert.assertTrue(
+ subscription.toString(),
subscription.getConsumerIds().contains(consumerId));
+ });
+ }
+ }
+
+ private static void awaitSubscriptionRemoved(final String topicName) throws
Exception {
+ try (final ISubscriptionTableSession session =
createSubscriptionSession()) {
+ Awaitility.await()
+ .pollInSameThread()
+ .pollInterval(Duration.ofMillis(500))
+ .atMost(Duration.ofSeconds(45))
+ .untilAsserted(
+ () -> {
+ final Set<Subscription> subscriptions =
session.getSubscriptions(topicName);
+ Assert.assertTrue(subscriptions.toString(),
subscriptions.isEmpty());
+ });
+ }
+ }
+
+ private static ISubscriptionTableSession createSubscriptionSession() throws
Exception {
+ final ISubscriptionTableSession session =
+ new SubscriptionTableSessionBuilder()
+ .host(EnvFactory.getEnv().getIP())
+ .port(Integer.parseInt(EnvFactory.getEnv().getPort()))
+ .build();
+ session.open();
+ return session;
+ }
+
+ private static void assertRowsConsumed(
+ final SubscriptionTablePullConsumer consumer, final Set<String>
expectedRows)
+ throws Exception {
+ final ConsensusSubscriptionTableITSupport.ConsumedRecords consumed =
+
ConsensusSubscriptionTableITSupport.pollAndCommitUntilContains(consumer,
expectedRows, 60);
+ ConsensusSubscriptionTableITSupport.assertExactRowKeys(expectedRows,
consumed);
+ }
+
+ private static Object getField(final Object target, final String fieldName)
throws Exception {
+ Class<?> currentClass = target.getClass();
+ while (currentClass != null) {
+ try {
+ final Field field = currentClass.getDeclaredField(fieldName);
+ field.setAccessible(true);
+ return field.get(target);
+ } catch (final NoSuchFieldException ignored) {
+ currentClass = currentClass.getSuperclass();
+ }
+ }
+ throw new NoSuchFieldException(fieldName);
+ }
+
+ private static void closeQuietly(final SubscriptionTablePullConsumer
consumer) {
+ if (consumer != null) {
+ try {
+ consumer.close();
+ } catch (final Exception ignored) {
+ // ignored on cleanup
+ }
+ }
+ }
+}
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 192a07dc83a..d918b4aed47 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
@@ -28,6 +28,8 @@ import
org.apache.iotdb.db.subscription.receiver.SubscriptionReceiver;
import org.apache.iotdb.db.subscription.receiver.SubscriptionReceiverV1;
import org.apache.iotdb.rpc.RpcUtils;
import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.rpc.subscription.config.ConsumerConfig;
+import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeRequestType;
import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeRequestVersion;
import
org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeResponseType;
import
org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeResponseVersion;
@@ -37,9 +39,10 @@ import
org.apache.iotdb.service.rpc.thrift.TPipeSubscribeResp;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Map;
-import java.util.Set;
+import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
@@ -49,9 +52,6 @@ public class SubscriptionReceiverAgent {
private static final Logger LOGGER =
LoggerFactory.getLogger(SubscriptionReceiverAgent.class);
- private static final Map<Byte, Supplier<SubscriptionReceiver>>
RECEIVER_CONSTRUCTORS =
- new HashMap<>();
-
private static final TPipeSubscribeResp SUBSCRIPTION_NOT_ENABLED_ERROR_RESP =
new TPipeSubscribeResp(
RpcUtils.getStatus(
@@ -60,21 +60,44 @@ public class SubscriptionReceiverAgent {
PipeSubscribeResponseVersion.VERSION_1.getVersion(),
PipeSubscribeResponseType.ACK.getType());
+ private final Map<Byte, Supplier<SubscriptionReceiver>> receiverConstructors
= new HashMap<>();
private final ThreadLocal<SubscriptionReceiver> receiverThreadLocal = new
ThreadLocal<>();
- private final Set<SubscriptionReceiver> activeReceivers =
ConcurrentHashMap.newKeySet();
- private final ScheduledExecutorService receiverTimeoutChecker =
- IoTDBThreadPoolFactory.newSingleThreadScheduledExecutor(
- SubscriptionReceiverAgent.class.getSimpleName() +
"-Timeout-Checker");
+
+ /**
+ * The receiver currently serving each consumer identity. A disconnected
receiver deliberately
+ * remains in this map until its inactivity timeout closes the consumer,
while a reconnecting
+ * receiver replaces it atomically through {@link
ConcurrentHashMap#compute(Object,
+ * java.util.function.BiFunction)}.
+ */
+ private final ConcurrentHashMap<ConsumerIdentity, SubscriptionReceiver>
consumerReceivers =
+ new ConcurrentHashMap<>();
+
+ private final ScheduledExecutorService receiverTimeoutChecker;
SubscriptionReceiverAgent() {
- RECEIVER_CONSTRUCTORS.put(
- PipeSubscribeRequestVersion.VERSION_1.getVersion(),
SubscriptionReceiverV1::new);
- ScheduledExecutorUtil.safelyScheduleWithFixedDelay(
- receiverTimeoutChecker,
- this::checkReceiverTimeouts,
- Math.max(1_000L,
SubscriptionConfig.getInstance().getSubscriptionDefaultTimeoutInMs() / 2L),
- Math.max(1_000L,
SubscriptionConfig.getInstance().getSubscriptionDefaultTimeoutInMs() / 2L),
- TimeUnit.MILLISECONDS);
+ this(SubscriptionReceiverV1::new, true);
+ }
+
+ SubscriptionReceiverAgent(
+ final Supplier<SubscriptionReceiver> receiverConstructor,
+ final boolean scheduleTimeoutChecker) {
+ receiverConstructors.put(
+ PipeSubscribeRequestVersion.VERSION_1.getVersion(),
receiverConstructor);
+ if (scheduleTimeoutChecker) {
+ receiverTimeoutChecker =
+ IoTDBThreadPoolFactory.newSingleThreadScheduledExecutor(
+ SubscriptionReceiverAgent.class.getSimpleName() +
"-Timeout-Checker");
+ ScheduledExecutorUtil.safelyScheduleWithFixedDelay(
+ receiverTimeoutChecker,
+ this::checkReceiverTimeouts,
+ Math.max(
+ 1_000L,
SubscriptionConfig.getInstance().getSubscriptionDefaultTimeoutInMs() / 2L),
+ Math.max(
+ 1_000L,
SubscriptionConfig.getInstance().getSubscriptionDefaultTimeoutInMs() / 2L),
+ TimeUnit.MILLISECONDS);
+ } else {
+ receiverTimeoutChecker = null;
+ }
}
public TPipeSubscribeResp handle(final TPipeSubscribeReq req) {
@@ -93,12 +116,48 @@ public class SubscriptionReceiverAgent {
}
final byte reqVersion = req.getVersion();
- if (RECEIVER_CONSTRUCTORS.containsKey(reqVersion)) {
+ if (receiverConstructors.containsKey(reqVersion)) {
final SubscriptionReceiver receiver = getReceiver(reqVersion);
receiver.setAuthenticatedUsername(username);
- activeReceivers.add(receiver);
- receiver.handleTimeout();
- return receiver.handle(req);
+ final ConsumerIdentity consumerIdentity = getConsumerIdentity(req,
receiver);
+ final RequestResult requestResult = new RequestResult();
+
+ if (Objects.isNull(consumerIdentity)) {
+ requestResult.response = handleRequest(receiver, req, null);
+ } else {
+ consumerReceivers.compute(
+ consumerIdentity,
+ (identity, currentReceiver) -> {
+ requestResult.response = handleRequest(receiver, req,
currentReceiver);
+
+ if (isHandshake(req)) {
+ if (isSuccessful(requestResult.response)) {
+ if (currentReceiver != null && currentReceiver != receiver) {
+ currentReceiver.invalidateConsumer();
+ }
+ return receiver;
+ }
+ return currentReceiver;
+ }
+
+ if (currentReceiver != null && currentReceiver != receiver) {
+ return currentReceiver;
+ }
+ return receiver.hasActiveConsumer() ? receiver : null;
+ });
+ }
+
+ if (isHandshake(req) && isSuccessful(requestResult.response)) {
+ final ConsumerIdentity activeIdentity = getConsumerIdentity(receiver);
+ if (!Objects.equals(consumerIdentity, activeIdentity)) {
+ registerReceiver(receiver, activeIdentity);
+ } else {
+ removeReceiverMappingsExcept(receiver, activeIdentity);
+ }
+ } else if (isClose(req) && isSuccessful(requestResult.response)) {
+ removeReceiverMappings(receiver);
+ }
+ return requestResult.response;
} else {
final TSStatus status =
RpcUtils.getStatus(
@@ -120,7 +179,7 @@ public class SubscriptionReceiverAgent {
}
public long remainingMs(final byte reqVersion) {
- if (RECEIVER_CONSTRUCTORS.containsKey(reqVersion)) {
+ if (receiverConstructors.containsKey(reqVersion)) {
return getReceiver(reqVersion).remainingMs();
} else {
return
SubscriptionConfig.getInstance().getSubscriptionDefaultTimeoutInMs();
@@ -147,8 +206,8 @@ public class SubscriptionReceiverAgent {
}
private SubscriptionReceiver setAndGetReceiver(final byte reqVersion) {
- if (RECEIVER_CONSTRUCTORS.containsKey(reqVersion)) {
- receiverThreadLocal.set(RECEIVER_CONSTRUCTORS.get(reqVersion).get());
+ if (receiverConstructors.containsKey(reqVersion)) {
+ receiverThreadLocal.set(receiverConstructors.get(reqVersion).get());
} else {
throw new UnsupportedOperationException(
String.format(
@@ -162,13 +221,140 @@ public class SubscriptionReceiverAgent {
public final void handleClientExit() {
final SubscriptionReceiver receiver = receiverThreadLocal.get();
if (receiver != null) {
- activeReceivers.remove(receiver);
- receiver.handleExit();
- receiverThreadLocal.remove();
+ try {
+ final ConsumerIdentity consumerIdentity =
getConsumerIdentity(receiver);
+ if (Objects.isNull(consumerIdentity)) {
+ receiver.handleExit();
+ } else {
+ consumerReceivers.compute(
+ consumerIdentity,
+ (identity, currentReceiver) -> {
+ if (currentReceiver != null && currentReceiver != receiver) {
+ // A newer connection has already taken over this consumer.
Do not let the old
+ // connection's exit cleanup touch the new owner's
subscription state.
+ receiver.invalidateConsumer();
+ receiver.handleExit();
+ return currentReceiver;
+ }
+ receiver.handleExit();
+ return receiver.hasActiveConsumer() ? receiver : null;
+ });
+ }
+ } finally {
+ receiverThreadLocal.remove();
+ }
+ }
+ }
+
+ void checkReceiverTimeouts() {
+ consumerReceivers.forEach(
+ (identity, receiver) ->
+ consumerReceivers.computeIfPresent(
+ identity,
+ (currentIdentity, currentReceiver) -> {
+ if (currentReceiver != receiver) {
+ return currentReceiver;
+ }
+ if (!identity.equals(getConsumerIdentity(receiver))) {
+ return null;
+ }
+ receiver.handleTimeout();
+ return receiver.hasActiveConsumer() ? receiver : null;
+ }));
+ }
+
+ private TPipeSubscribeResp handleRequest(
+ final SubscriptionReceiver receiver,
+ final TPipeSubscribeReq req,
+ final SubscriptionReceiver currentReceiver) {
+ if (!isHandshake(req) && currentReceiver != null && currentReceiver !=
receiver) {
+ receiver.invalidateConsumer();
+ }
+ return receiver.handle(req);
+ }
+
+ private void registerReceiver(
+ final SubscriptionReceiver receiver, final ConsumerIdentity identity) {
+ if (Objects.isNull(identity)) {
+ removeReceiverMappings(receiver);
+ return;
+ }
+ consumerReceivers.compute(
+ identity,
+ (key, currentReceiver) -> {
+ if (currentReceiver != null && currentReceiver != receiver) {
+ currentReceiver.invalidateConsumer();
+ }
+ return receiver;
+ });
+ removeReceiverMappingsExcept(receiver, identity);
+ }
+
+ private void removeReceiverMappingsExcept(
+ final SubscriptionReceiver receiver, final ConsumerIdentity
retainedIdentity) {
+ consumerReceivers.forEach(
+ (registeredIdentity, currentReceiver) -> {
+ if (currentReceiver == receiver
+ && !Objects.equals(retainedIdentity, registeredIdentity)) {
+ consumerReceivers.remove(registeredIdentity, receiver);
+ }
+ });
+ }
+
+ private void removeReceiverMappings(final SubscriptionReceiver receiver) {
+ consumerReceivers.forEach(
+ (identity, currentReceiver) -> consumerReceivers.remove(identity,
receiver));
+ }
+
+ private static ConsumerIdentity getConsumerIdentity(
+ final TPipeSubscribeReq req, final SubscriptionReceiver receiver) {
+ if (isHandshake(req) && req.isSetBody()) {
+ try {
+ final ByteBuffer body = req.bufferForBody();
+ if (body.hasRemaining()) {
+ final ConsumerConfig consumerConfig =
ConsumerConfig.deserialize(body);
+ final ConsumerIdentity identity =
+ ConsumerIdentity.of(
+ consumerConfig.getConsumerGroupId(),
consumerConfig.getConsumerId());
+ if (Objects.nonNull(identity)) {
+ return identity;
+ }
+ }
+ } catch (final RuntimeException ignored) {
+ // Let the receiver report the malformed handshake request. It still
needs to see the
+ // original buffer, so parsing is intentionally done on a duplicate
above.
+ }
}
+ return getConsumerIdentity(receiver);
+ }
+
+ private static ConsumerIdentity getConsumerIdentity(final
SubscriptionReceiver receiver) {
+ return ConsumerIdentity.of(receiver.getConsumerGroupId(),
receiver.getConsumerId());
}
- private void checkReceiverTimeouts() {
- activeReceivers.forEach(SubscriptionReceiver::handleTimeout);
+ private static boolean isHandshake(final TPipeSubscribeReq req) {
+ return req.getType() == PipeSubscribeRequestType.HANDSHAKE.getType();
+ }
+
+ private static boolean isClose(final TPipeSubscribeReq req) {
+ return req.getType() == PipeSubscribeRequestType.CLOSE.getType();
+ }
+
+ private static boolean isSuccessful(final TPipeSubscribeResp response) {
+ return response != null
+ && response.getStatus() != null
+ && response.getStatus().getCode() ==
TSStatusCode.SUCCESS_STATUS.getStatusCode();
+ }
+
+ private static final class RequestResult {
+ private TPipeSubscribeResp response;
+ }
+
+ private record ConsumerIdentity(String consumerGroupId, String consumerId) {
+ private static ConsumerIdentity of(final String consumerGroupId, final
String consumerId) {
+ return Objects.isNull(consumerGroupId) || Objects.isNull(consumerId)
+ ? null
+ : new ConsumerIdentity(consumerGroupId, consumerId);
+ }
}
}
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 cc7b57eee81..e5de617c89c 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
@@ -35,5 +35,26 @@ public interface SubscriptionReceiver {
void handleTimeout();
+ /**
+ * Returns the identity of the consumer currently associated with this
receiver, or {@code null}
+ * if the receiver has not completed a handshake (or has already been
invalidated).
+ */
+ String getConsumerId();
+
+ /**
+ * Returns the consumer group currently associated with this receiver, or
{@code null} if the
+ * receiver has not completed a handshake (or has already been invalidated).
+ */
+ String getConsumerGroupId();
+
+ /**
+ * Invalidates this receiver so that requests from an obsolete connection
cannot affect a new
+ * owner.
+ */
+ void invalidateConsumer();
+
+ /** Returns whether this receiver still owns an active consumer timeout
state. */
+ boolean hasActiveConsumer();
+
long remainingMs();
}
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 4ca93fb5cb8..dcc92bc2217 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
@@ -128,6 +128,7 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
private volatile boolean consumerInvalidated;
private volatile long lastActivityTimeMs = System.currentTimeMillis();
private final AtomicLong inFlightRequestCount = new AtomicLong(0);
+ private long consumerStateVersion;
private static final String SQL_DIALECT_TABLE_VALUE = "table";
@@ -186,59 +187,91 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
authenticatedUsername = username;
}
+ @Override
+ public String getConsumerId() {
+ final ConsumerConfig consumerConfig = sharedConsumerConfig;
+ return Objects.isNull(consumerConfig) ? null :
consumerConfig.getConsumerId();
+ }
+
+ @Override
+ public String getConsumerGroupId() {
+ final ConsumerConfig consumerConfig = sharedConsumerConfig;
+ return Objects.isNull(consumerConfig) ? null :
consumerConfig.getConsumerGroupId();
+ }
+
+ @Override
+ public boolean hasActiveConsumer() {
+ return Objects.nonNull(sharedConsumerConfig);
+ }
+
+ @Override
+ public void invalidateConsumer() {
+ clearSharedConsumerState();
+ }
+
@Override
public void handleExit() {
- final ConsumerConfig consumerConfig = consumerConfigThreadLocal.get();
- if (Objects.nonNull(consumerConfig)) {
- LOGGER.info(
- DataNodePipeMessages
-
.PIPE_LOG_SUBSCRIPTION_REMOVE_CONSUMER_CONFIG_WHEN_HANDLING_EXIT_3827D0E8,
- consumerConfigThreadLocal.get());
- // we should not close the consumer here because it might reuse the
previous consumption
- // progress to continue consuming
- // closeConsumer(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
- unsubscribeCompleteTopics(consumerConfig);
- consumerConfigThreadLocal.remove();
+ synchronized (this) {
+ final ConsumerConfig consumerConfig = consumerConfigThreadLocal.get();
+ try {
+ if (Objects.nonNull(consumerConfig)) {
+ LOGGER.info(
+ DataNodePipeMessages
+
.PIPE_LOG_SUBSCRIPTION_REMOVE_CONSUMER_CONFIG_WHEN_HANDLING_EXIT_3827D0E8,
+ consumerConfig);
+ // We should not close the consumer here because it might reuse the
previous consumption
+ // 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) {
+ // 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.
+ unsubscribeCompleteTopics(consumerConfig);
+ }
+ }
+ } finally {
+ consumerConfigThreadLocal.remove();
+ pollTimerThreadLocal.remove();
+ authenticatedUsername = null;
+ }
}
- clearSharedConsumerState();
- authenticatedUsername = null;
}
@Override
public void handleTimeout() {
- final ConsumerConfig consumerConfig;
- final long inactiveMs;
- final long timeoutMs;
synchronized (this) {
- consumerConfig = sharedConsumerConfig;
+ final ConsumerConfig consumerConfig = sharedConsumerConfig;
if (Objects.isNull(consumerConfig) || inFlightRequestCount.get() > 0) {
return;
}
- timeoutMs = calculateConsumerInactivityTimeoutMs(consumerConfig);
- inactiveMs = System.currentTimeMillis() - lastActivityTimeMs;
+ final long timeoutMs =
calculateConsumerInactivityTimeoutMs(consumerConfig);
+ final long inactiveMs = System.currentTimeMillis() - lastActivityTimeMs;
if (inactiveMs <= timeoutMs) {
return;
}
- clearSharedConsumerState();
- }
- LOGGER.info(
- DataNodePipeMessages
-
.PIPE_LOG_SUBSCRIPTION_CONSUMER_IS_INACTIVE_FOR_MS_EXCEEDING_TIMEOUT_36E06B11,
- consumerConfig,
- inactiveMs,
- timeoutMs);
- try {
- closeConsumer(consumerConfig);
- } catch (final Exception e) {
- LOGGER.warn(
+ final long stateVersion = consumerStateVersion;
+ LOGGER.info(
DataNodePipeMessages
-
.PIPE_LOG_SUBSCRIPTION_FAILED_TO_CLOSE_TIMED_OUT_CONSUMER_AFTER_MS_89CC11F1,
+
.PIPE_LOG_SUBSCRIPTION_CONSUMER_IS_INACTIVE_FOR_MS_EXCEEDING_TIMEOUT_36E06B11,
consumerConfig,
inactiveMs,
- e);
+ timeoutMs);
+ try {
+ closeConsumer(consumerConfig);
+ // Clear the timeout state only after successful cleanup and only if
it still represents the
+ // same consumer activation. On failure, keep it so the timeout
checker can retry.
+ if (sharedConsumerConfig == consumerConfig && consumerStateVersion ==
stateVersion) {
+ clearSharedConsumerState();
+ }
+ } catch (final Exception e) {
+ LOGGER.warn(
+ DataNodePipeMessages
+
.PIPE_LOG_SUBSCRIPTION_FAILED_TO_CLOSE_TIMED_OUT_CONSUMER_AFTER_MS_89CC11F1,
+ consumerConfig,
+ inactiveMs,
+ e);
+ }
}
}
@@ -1084,7 +1117,7 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
return PipeSubscribeSeekResp.toTPipeSubscribeResp(RpcUtils.SUCCESS_STATUS);
}
- private void closeConsumer(final ConsumerConfig consumerConfig) {
+ void closeConsumer(final ConsumerConfig consumerConfig) {
// unsubscribe all subscribed topics
final Set<String> topicNames =
SubscriptionAgent.consumer()
@@ -1318,12 +1351,16 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
sharedConsumerConfig = consumerConfig;
consumerInvalidated = false;
lastActivityTimeMs = System.currentTimeMillis();
+ consumerStateVersion++;
}
}
- private void clearSharedConsumerState() {
- sharedConsumerConfig = null;
- consumerInvalidated = true;
+ private synchronized void clearSharedConsumerState() {
+ if (Objects.nonNull(sharedConsumerConfig) || !consumerInvalidated) {
+ sharedConsumerConfig = null;
+ consumerInvalidated = true;
+ consumerStateVersion++;
+ }
}
private long calculateConsumerInactivityTimeoutMs(final ConsumerConfig
consumerConfig) {
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
new file mode 100644
index 00000000000..a151e83e21d
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/agent/SubscriptionReceiverAgentTest.java
@@ -0,0 +1,273 @@
+/*
+ * 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.iotdb.db.subscription.agent;
+
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.db.subscription.receiver.SubscriptionReceiver;
+import org.apache.iotdb.rpc.RpcUtils;
+import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.rpc.subscription.config.ConsumerConfig;
+import org.apache.iotdb.rpc.subscription.config.ConsumerConstant;
+import org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeCloseReq;
+import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeHandshakeReq;
+import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeRequestType;
+import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeRequestVersion;
+import
org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeResponseType;
+import
org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeResponseVersion;
+import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeReq;
+import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeResp;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
+
+public class SubscriptionReceiverAgentTest {
+
+ @Test
+ public void testDisconnectedReceiverIsRetainedUntilTimeout() throws
IOException {
+ final CopyOnWriteArrayList<FakeSubscriptionReceiver> receivers = new
CopyOnWriteArrayList<>();
+ final SubscriptionReceiverAgent agent = createAgent(receivers, true /*
closeOnTimeout */);
+ final TPipeSubscribeReq handshake = createHandshakeRequest("group",
"consumer");
+
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ agent.handle(handshake, "root").getStatus().getCode());
+ agent.handleClientExit();
+
+ final FakeSubscriptionReceiver receiver = receivers.get(0);
+ Assert.assertEquals(1, receiver.exitCount.get());
+ Assert.assertEquals(0, receiver.timeoutCount.get());
+ agent.checkReceiverTimeouts();
+
+ Assert.assertEquals(1, receiver.timeoutCount.get());
+ Assert.assertFalse(receiver.hasActiveConsumer());
+ agent.checkReceiverTimeouts();
+ Assert.assertEquals(1, receiver.timeoutCount.get());
+ }
+
+ @Test
+ public void testReconnectInvalidatesOldReceiverBeforeTimeoutCleanup() throws
IOException {
+ final CopyOnWriteArrayList<FakeSubscriptionReceiver> receivers = new
CopyOnWriteArrayList<>();
+ final SubscriptionReceiverAgent agent = createAgent(receivers, false /*
closeOnTimeout */);
+ final TPipeSubscribeReq handshake = createHandshakeRequest("group",
"consumer");
+
+ agent.handle(handshake, "root");
+ agent.handleClientExit();
+ agent.handle(handshake, "root");
+
+ final FakeSubscriptionReceiver oldReceiver = receivers.get(0);
+ final FakeSubscriptionReceiver newReceiver = receivers.get(1);
+ Assert.assertTrue(oldReceiver.invalidated);
+ Assert.assertFalse(newReceiver.invalidated);
+
+ agent.checkReceiverTimeouts();
+
+ Assert.assertEquals(0, oldReceiver.timeoutCount.get());
+ Assert.assertEquals(1, newReceiver.timeoutCount.get());
+ }
+
+ @Test
+ public void testLateExitFromOldConnectionKeepsNewReceiverRegistered() throws
Exception {
+ final CopyOnWriteArrayList<FakeSubscriptionReceiver> receivers = new
CopyOnWriteArrayList<>();
+ final SubscriptionReceiverAgent agent = createAgent(receivers, false /*
closeOnTimeout */);
+ final CountDownLatch oldHandshakeCompleted = new CountDownLatch(1);
+ final CountDownLatch allowOldConnectionToExit = new CountDownLatch(1);
+ final AtomicReference<Throwable> threadFailure = new AtomicReference<>();
+
+ final Thread oldConnection =
+ new Thread(
+ () -> {
+ try {
+ agent.handle(createHandshakeRequest("group", "consumer"),
"root");
+ oldHandshakeCompleted.countDown();
+ allowOldConnectionToExit.await();
+ agent.handleClientExit();
+ } catch (final Throwable t) {
+ threadFailure.set(t);
+ oldHandshakeCompleted.countDown();
+ }
+ });
+ oldConnection.start();
+ final Thread newConnection =
+ new Thread(
+ () -> {
+ try {
+ agent.handle(createHandshakeRequest("group", "consumer"),
"root");
+ } catch (final Throwable t) {
+ threadFailure.set(t);
+ }
+ });
+ try {
+ Assert.assertTrue(oldHandshakeCompleted.await(10, TimeUnit.SECONDS));
+ newConnection.start();
+ newConnection.join(TimeUnit.SECONDS.toMillis(10));
+ Assert.assertFalse(newConnection.isAlive());
+ } finally {
+ allowOldConnectionToExit.countDown();
+ oldConnection.join(TimeUnit.SECONDS.toMillis(10));
+ }
+ Assert.assertFalse(oldConnection.isAlive());
+
+ if (threadFailure.get() != null) {
+ throw new AssertionError(threadFailure.get());
+ }
+ final FakeSubscriptionReceiver oldReceiver = receivers.get(0);
+ final FakeSubscriptionReceiver newReceiver = receivers.get(1);
+ agent.checkReceiverTimeouts();
+
+ Assert.assertTrue(oldReceiver.invalidated);
+ Assert.assertEquals(0, oldReceiver.timeoutCount.get());
+ Assert.assertEquals(1, newReceiver.timeoutCount.get());
+ }
+
+ @Test
+ public void testSuccessfulCloseRemovesReceiverFromTimeoutRegistry() throws
IOException {
+ final CopyOnWriteArrayList<FakeSubscriptionReceiver> receivers = new
CopyOnWriteArrayList<>();
+ final SubscriptionReceiverAgent agent = createAgent(receivers, true /*
closeOnTimeout */);
+ final TPipeSubscribeReq handshake = createHandshakeRequest("group",
"consumer");
+
+ agent.handle(handshake, "root");
+ final TPipeSubscribeResp closeResponse =
+ agent.handle(PipeSubscribeCloseReq.toTPipeSubscribeReq(), "root");
+
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
closeResponse.getStatus().getCode());
+ final FakeSubscriptionReceiver receiver = receivers.get(0);
+ Assert.assertEquals(0, receiver.timeoutCount.get());
+ agent.checkReceiverTimeouts();
+ Assert.assertEquals(0, receiver.timeoutCount.get());
+ }
+
+ private SubscriptionReceiverAgent createAgent(
+ final CopyOnWriteArrayList<FakeSubscriptionReceiver> receivers,
+ final boolean closeOnTimeout) {
+ final Supplier<SubscriptionReceiver> constructor =
+ () -> {
+ final FakeSubscriptionReceiver receiver = new
FakeSubscriptionReceiver(closeOnTimeout);
+ receivers.add(receiver);
+ return receiver;
+ };
+ return new SubscriptionReceiverAgent(constructor, false);
+ }
+
+ private TPipeSubscribeReq createHandshakeRequest(
+ final String consumerGroupId, final String consumerId) throws
IOException {
+ final Map<String, String> attributes = new HashMap<>();
+ attributes.put(ConsumerConstant.CONSUMER_GROUP_ID_KEY, consumerGroupId);
+ attributes.put(ConsumerConstant.CONSUMER_ID_KEY, consumerId);
+ return PipeSubscribeHandshakeReq.toTPipeSubscribeReq(new
ConsumerConfig(attributes));
+ }
+
+ private static class FakeSubscriptionReceiver implements
SubscriptionReceiver {
+
+ private final boolean closeOnTimeout;
+ private final AtomicInteger timeoutCount = new AtomicInteger();
+ private final AtomicInteger exitCount = new AtomicInteger();
+ private ConsumerConfig consumerConfig;
+ private boolean invalidated;
+
+ private FakeSubscriptionReceiver(final boolean closeOnTimeout) {
+ this.closeOnTimeout = closeOnTimeout;
+ }
+
+ @Override
+ public TPipeSubscribeResp handle(final TPipeSubscribeReq req) {
+ if (req.getType() == PipeSubscribeRequestType.HANDSHAKE.getType()) {
+ consumerConfig = ConsumerConfig.deserialize(req.bufferForBody());
+ invalidated = false;
+ return response(TSStatusCode.SUCCESS_STATUS);
+ }
+ if (req.getType() == PipeSubscribeRequestType.CLOSE.getType()) {
+ consumerConfig = null;
+ invalidated = true;
+ return response(TSStatusCode.SUCCESS_STATUS);
+ }
+ return response(
+ invalidated ? TSStatusCode.SUBSCRIPTION_MISSING_CONSUMER :
TSStatusCode.SUCCESS_STATUS);
+ }
+
+ @Override
+ public void setAuthenticatedUsername(final String username) {
+ // no-op
+ }
+
+ @Override
+ public PipeSubscribeRequestVersion getVersion() {
+ return PipeSubscribeRequestVersion.VERSION_1;
+ }
+
+ @Override
+ public void handleExit() {
+ exitCount.incrementAndGet();
+ }
+
+ @Override
+ public void handleTimeout() {
+ timeoutCount.incrementAndGet();
+ if (closeOnTimeout) {
+ consumerConfig = null;
+ invalidated = true;
+ }
+ }
+
+ @Override
+ public String getConsumerId() {
+ return consumerConfig == null ? null : consumerConfig.getConsumerId();
+ }
+
+ @Override
+ public String getConsumerGroupId() {
+ return consumerConfig == null ? null :
consumerConfig.getConsumerGroupId();
+ }
+
+ @Override
+ public void invalidateConsumer() {
+ consumerConfig = null;
+ invalidated = true;
+ }
+
+ @Override
+ public boolean hasActiveConsumer() {
+ return consumerConfig != null;
+ }
+
+ @Override
+ public long remainingMs() {
+ return 0;
+ }
+
+ private TPipeSubscribeResp response(final TSStatusCode statusCode) {
+ final TSStatus status = RpcUtils.getStatus(statusCode);
+ return new TPipeSubscribeResp(
+ status,
+ PipeSubscribeResponseVersion.VERSION_1.getVersion(),
+ PipeSubscribeResponseType.ACK.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 0890ae8100a..21d2e7b67d8 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
@@ -69,6 +69,96 @@ public class SubscriptionReceiverV1Test {
Assert.assertFalse((boolean) getField(receiver, "consumerInvalidated"));
}
+ @Test
+ public void testHandleTimeoutRetriesAfterInFlightRequestCompletes() throws
Exception {
+ final AtomicLong closeAttemptCount = new AtomicLong();
+ final SubscriptionReceiverV1 receiver =
+ new SubscriptionReceiverV1() {
+ @Override
+ void closeConsumer(final ConsumerConfig consumerConfig) {
+ closeAttemptCount.incrementAndGet();
+ }
+ };
+ final ConsumerConfig consumerConfig = createConsumerConfig(1_000L);
+ setField(receiver, "sharedConsumerConfig", consumerConfig);
+ final long timeoutMs =
invokeCalculateConsumerInactivityTimeoutMs(receiver, consumerConfig);
+ setField(receiver, "lastActivityTimeMs", System.currentTimeMillis() -
timeoutMs - 1L);
+ final AtomicLong inFlightRequestCount = (AtomicLong) getField(receiver,
"inFlightRequestCount");
+
+ inFlightRequestCount.set(1L);
+ receiver.handleTimeout();
+
+ Assert.assertEquals(0L, closeAttemptCount.get());
+ Assert.assertSame(consumerConfig, getField(receiver,
"sharedConsumerConfig"));
+
+ inFlightRequestCount.set(0L);
+ receiver.handleTimeout();
+
+ Assert.assertEquals(1L, closeAttemptCount.get());
+ Assert.assertNull(getField(receiver, "sharedConsumerConfig"));
+ Assert.assertTrue((boolean) getField(receiver, "consumerInvalidated"));
+ }
+
+ @Test
+ public void testHandleTimeoutRetriesAfterCleanupFailure() throws Exception {
+ final AtomicLong closeAttemptCount = new AtomicLong();
+ final SubscriptionReceiverV1 receiver =
+ new SubscriptionReceiverV1() {
+ @Override
+ void closeConsumer(final ConsumerConfig consumerConfig) {
+ closeAttemptCount.incrementAndGet();
+ throw new RuntimeException("expected cleanup failure");
+ }
+ };
+ final ConsumerConfig consumerConfig = createConsumerConfig(1_000L);
+ setField(receiver, "sharedConsumerConfig", consumerConfig);
+ final long timeoutMs =
invokeCalculateConsumerInactivityTimeoutMs(receiver, consumerConfig);
+ setField(receiver, "lastActivityTimeMs", System.currentTimeMillis() -
timeoutMs - 1L);
+
+ receiver.handleTimeout();
+ receiver.handleTimeout();
+
+ Assert.assertEquals(2L, closeAttemptCount.get());
+ Assert.assertSame(consumerConfig, getField(receiver,
"sharedConsumerConfig"));
+ Assert.assertFalse((boolean) getField(receiver, "consumerInvalidated"));
+ }
+
+ @Test
+ @SuppressWarnings("unchecked")
+ public void testHandleExitKeepsSharedConsumerStateForTimeoutCleanup() throws
Exception {
+ final SubscriptionReceiverV1 receiver = new SubscriptionReceiverV1();
+ final ConsumerConfig consumerConfig = createConsumerConfig(1_000L);
+ setField(receiver, "sharedConsumerConfig", consumerConfig);
+ final ThreadLocal<ConsumerConfig> consumerConfigThreadLocal =
+ (ThreadLocal<ConsumerConfig>) getField(receiver,
"consumerConfigThreadLocal");
+ consumerConfigThreadLocal.set(consumerConfig);
+
+ receiver.handleExit();
+
+ Assert.assertSame(consumerConfig, getField(receiver,
"sharedConsumerConfig"));
+ Assert.assertFalse((boolean) getField(receiver, "consumerInvalidated"));
+ Assert.assertTrue(receiver.hasActiveConsumer());
+ Assert.assertNull(consumerConfigThreadLocal.get());
+ }
+
+ @Test
+ @SuppressWarnings("unchecked")
+ public void testHandleExitClearsThreadLocalStateAfterInvalidation() throws
Exception {
+ final SubscriptionReceiverV1 receiver = new SubscriptionReceiverV1();
+ final ConsumerConfig consumerConfig = createConsumerConfig(1_000L);
+ setField(receiver, "sharedConsumerConfig", consumerConfig);
+ final ThreadLocal<ConsumerConfig> consumerConfigThreadLocal =
+ (ThreadLocal<ConsumerConfig>) getField(receiver,
"consumerConfigThreadLocal");
+ consumerConfigThreadLocal.set(consumerConfig);
+
+ receiver.invalidateConsumer();
+ receiver.handleExit();
+
+ Assert.assertFalse(receiver.hasActiveConsumer());
+ Assert.assertTrue((boolean) getField(receiver, "consumerInvalidated"));
+ Assert.assertNull(consumerConfigThreadLocal.get());
+ }
+
@Test
public void testCalculateConsumerInactivityTimeoutUsesDefaultTimeout()
throws Exception {
final SubscriptionReceiverV1 receiver = new SubscriptionReceiverV1();
@@ -195,14 +285,14 @@ public class SubscriptionReceiverV1Test {
}
private Object getField(final Object target, final String fieldName) throws
Exception {
- final Field field = target.getClass().getDeclaredField(fieldName);
+ final Field field =
SubscriptionReceiverV1.class.getDeclaredField(fieldName);
field.setAccessible(true);
return field.get(target);
}
private void setField(final Object target, final String fieldName, final
Object value)
throws Exception {
- final Field field = target.getClass().getDeclaredField(fieldName);
+ final Field field =
SubscriptionReceiverV1.class.getDeclaredField(fieldName);
field.setAccessible(true);
field.set(target, value);
}
