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 6709dfa9494 Fix concurrent subscription consumer handshakes (#18654)
6709dfa9494 is described below
commit 6709dfa9494cb7517ad8b940aa301fe93a03b43b
Author: Caideyipi <[email protected]>
AuthorDate: Fri Sep 18 17:37:37 2026 +0800
Fix concurrent subscription consumer handshakes (#18654)
* Fix concurrent subscription consumer handshakes
* Make consumer fencing consistent across DataNodes
---
.../java/org/apache/iotdb/rpc/TSStatusCode.java | 1 +
.../rpc/subscription/config/ConsumerConfig.java | 4 +
.../rpc/subscription/config/ConsumerConstant.java | 1 +
.../SubscriptionConsumerFencedException.java | 51 +++
.../base/AbstractSubscriptionConsumer.java | 200 +++++++++---
.../base/AbstractSubscriptionProvider.java | 21 ++
.../base/AbstractSubscriptionProviders.java | 55 +++-
.../base/AbstractSubscriptionPullConsumer.java | 14 +-
.../base/AbstractSubscriptionPushConsumer.java | 4 +-
.../base/SubscriptionConsumerLifecycleTest.java | 352 +++++++++++++++++---
.../base/SubscriptionProviderStatusTest.java | 72 ++++
.../apache/iotdb/db/i18n/DataNodePipeMessages.java | 6 +
.../apache/iotdb/db/i18n/DataNodePipeMessages.java | 6 +
.../agent/SubscriptionReceiverAgent.java | 86 ++++-
.../receiver/SubscriptionReceiver.java | 8 +
.../receiver/SubscriptionReceiverV1.java | 42 ++-
.../agent/SubscriptionReceiverAgentTest.java | 361 ++++++++++++++++++++-
.../receiver/SubscriptionReceiverV1Test.java | 41 +++
18 files changed, 1197 insertions(+), 128 deletions(-)
diff --git
a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSStatusCode.java
b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSStatusCode.java
index 61431249eb3..fce965fe27b 100644
---
a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSStatusCode.java
+++
b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSStatusCode.java
@@ -332,6 +332,7 @@ public enum TSStatusCode {
SUBSCRIPTION_OWNER_EPOCH_REQUIRED(1916),
SUBSCRIPTION_OWNER_LEASE_EXPIRED(1917),
SUBSCRIPTION_OWNER_EPOCH_CONFLICT(1918),
+ SUBSCRIPTION_CONSUMER_FENCED(1919),
// Topic
CREATE_TOPIC_ERROR(2000),
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConfig.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConfig.java
index 13f2a9ee3fb..47926b5fc91 100644
---
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConfig.java
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConfig.java
@@ -68,6 +68,10 @@ public class ConsumerConfig extends PipeParameters {
return getString(ConsumerConstant.CONSUMER_GROUP_ID_KEY);
}
+ public String getConsumerInstanceId() {
+ return getString(ConsumerConstant.CONSUMER_INSTANCE_ID_KEY);
+ }
+
public String getOwnerId() {
return getString(ConsumerConstant.OWNER_ID_KEY);
}
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConstant.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConstant.java
index 3df95facf36..89878efb49a 100644
---
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConstant.java
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConstant.java
@@ -40,6 +40,7 @@ public class ConsumerConstant {
public static final String CONSUMER_ID_KEY = "consumer-id";
public static final String CONSUMER_GROUP_ID_KEY = "group-id";
+ public static final String CONSUMER_INSTANCE_ID_KEY = "consumer-instance-id";
public static final String OWNER_ID_KEY = "owner-id";
public static final String OWNER_EPOCH_KEY = "owner-epoch";
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/exception/SubscriptionConsumerFencedException.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/exception/SubscriptionConsumerFencedException.java
new file mode 100644
index 00000000000..ef8e83b2c7c
--- /dev/null
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/exception/SubscriptionConsumerFencedException.java
@@ -0,0 +1,51 @@
+/*
+ * 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.rpc.subscription.exception;
+
+import java.util.Objects;
+
+/**
+ * Indicates that another connection has taken over the same subscription
consumer identity. The
+ * fenced consumer instance can no longer issue requests and must not reclaim
ownership
+ * automatically.
+ */
+public class SubscriptionConsumerFencedException extends
SubscriptionRuntimeCriticalException {
+
+ public SubscriptionConsumerFencedException(final String message) {
+ super(message);
+ }
+
+ public SubscriptionConsumerFencedException(final String message, final
Throwable cause) {
+ super(message, cause);
+ }
+
+ @Override
+ public boolean equals(final Object obj) {
+ return obj instanceof SubscriptionConsumerFencedException
+ && Objects.equals(getMessage(), ((SubscriptionConsumerFencedException)
obj).getMessage())
+ && Objects.equals(
+ getTimeStamp(), ((SubscriptionConsumerFencedException)
obj).getTimeStamp());
+ }
+
+ @Override
+ public int hashCode() {
+ return super.hashCode();
+ }
+}
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java
index 8a2c9ff6334..a4f5d36a324 100644
---
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java
@@ -24,6 +24,7 @@ import org.apache.iotdb.isession.SessionConfig;
import org.apache.iotdb.rpc.subscription.config.ConsumerConstant;
import org.apache.iotdb.rpc.subscription.config.TopicConfig;
import
org.apache.iotdb.rpc.subscription.exception.SubscriptionConnectionException;
+import
org.apache.iotdb.rpc.subscription.exception.SubscriptionConsumerFencedException;
import org.apache.iotdb.rpc.subscription.exception.SubscriptionException;
import
org.apache.iotdb.rpc.subscription.exception.SubscriptionOwnerFencedException;
import
org.apache.iotdb.rpc.subscription.exception.SubscriptionPipeTimeoutException;
@@ -90,6 +91,8 @@ import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiFunction;
import java.util.stream.Collectors;
@@ -107,6 +110,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
private static final long SLEEP_MS = 100L;
private static final long SLEEP_DELTA_MS = 50L;
private static final long TIMER_DELTA_MS = 250L;
+ private static final AtomicLong LAST_CONSUMER_INSTANCE_EPOCH = new
AtomicLong();
private final String username;
private final String password;
@@ -116,6 +120,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
protected String consumerGroupId;
protected String ownerId;
protected Long ownerEpoch;
+ private final String consumerInstanceId = generateConsumerInstanceId();
private final long heartbeatIntervalMs;
private final long endpointsSyncIntervalMs;
@@ -123,6 +128,8 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
private final AbstractSubscriptionProviders providers;
private final AtomicBoolean isClosed = new AtomicBoolean(true);
+ private final AtomicReference<SubscriptionConsumerFencedException>
fencedException =
+ new AtomicReference<>();
// This variable indicates whether the consumer has ever been closed.
private final AtomicBoolean isReleased = new AtomicBoolean(false);
@@ -193,6 +200,10 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
return ownerEpoch;
}
+ String getConsumerInstanceId() {
+ return consumerInstanceId;
+ }
+
/////////////////////////////// ctor ///////////////////////////////
protected AbstractSubscriptionConsumer(final
AbstractSubscriptionConsumerBuilder builder) {
@@ -296,6 +307,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
/////////////////////////////// open & close ///////////////////////////////
private void checkIfHasBeenClosed() throws SubscriptionException {
+ checkIfFenced();
if (isReleased.get()) {
final String errorMessage =
String.format("%s has ever been closed, unsupported operation after
closing.", this);
@@ -305,6 +317,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
}
private void checkIfOpened() throws SubscriptionException {
+ checkIfFenced();
if (isClosed.get()) {
final String errorMessage =
String.format("%s is not yet open, please open the subscription
consumer first.", this);
@@ -324,6 +337,9 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
providers.acquireWriteLock();
try {
providers.openProviders(this); // throw SubscriptionException
+ } catch (final SubscriptionException e) {
+ providers.closeProviders(!isFenced());
+ throw e;
} finally {
providers.releaseWriteLock();
}
@@ -354,13 +370,28 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
// Do not wait for the providers write lock here. A poll or heartbeat may
hold it while blocked
// in network I/O. prepareClose() bounds or interrupts those RPCs before
providers are detached.
providers.prepareClose();
- providers.closeProviders();
+ providers.closeProviders(!isFenced());
}
boolean isClosed() {
return isClosed.get();
}
+ boolean isFenced() {
+ return fencedException.get() != null;
+ }
+
+ void fence(final SubscriptionConsumerFencedException e) {
+ fencedException.compareAndSet(null, e);
+ }
+
+ void checkIfFenced() {
+ final SubscriptionConsumerFencedException e = fencedException.get();
+ if (e != null) {
+ throw e;
+ }
+ }
+
/////////////////////////////// subscribe & unsubscribe
///////////////////////////////
protected void subscribe(final String topicName) throws
SubscriptionException {
@@ -557,13 +588,18 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
this.thriftMaxFrameSize,
this.heartbeatIntervalMs,
this.connectionTimeoutInMs);
+ provider.setConsumerInstanceId(consumerInstanceId);
try {
provider.handshake();
} catch (final Exception e) {
try {
- provider.close();
+ provider.closeSession();
} catch (final Exception ignored) {
}
+ if (e instanceof SubscriptionConsumerFencedException) {
+ fence((SubscriptionConsumerFencedException) e);
+ throw (SubscriptionConsumerFencedException) e;
+ }
throw new SubscriptionConnectionException(
String.format(
SubscriptionMessages
@@ -598,6 +634,13 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
return message;
}
+ private static String generateConsumerInstanceId() {
+ final long epoch =
+ LAST_CONSUMER_INSTANCE_EPOCH.updateAndGet(
+ previous -> Math.max(System.currentTimeMillis(), previous + 1));
+ return String.format("%016x-%s", epoch,
RandomStringGenerator.generate(16));
+ }
+
/////////////////////////////// file ops ///////////////////////////////
private Path getFileDir(final String topicName) throws IOException {
@@ -734,6 +777,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
protected List<SubscriptionMessage> multiplePoll(
/* @NotNull */ final Set<String> topicNames, final long timeoutMs) {
+ checkIfFenced();
if (topicNames.isEmpty()) {
return Collections.emptyList();
}
@@ -754,42 +798,10 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
tasks.add(new PollTask(partition, timeoutMs));
}
- // submit multiple tasks to poll messages
- final List<SubscriptionMessage> messages = new ArrayList<>();
- SubscriptionRuntimeCriticalException
lastSubscriptionRuntimeCriticalException = null;
try {
// strict timeout
- for (final Future<List<SubscriptionMessage>> future :
- SubscriptionExecutorServiceManager.submitMultiplePollTasks(tasks,
timeoutMs)) {
- try {
- if (future.isCancelled()) {
- continue;
- }
- messages.addAll(future.get());
- } catch (final CancellationException ignored) {
-
- } catch (final ExecutionException e) {
- final Throwable cause = e.getCause();
- if (cause instanceof SubscriptionRuntimeCriticalException) {
- final SubscriptionRuntimeCriticalException ex =
- (SubscriptionRuntimeCriticalException) cause;
- LOGGER.warn(
- SubscriptionMessages
-
.LOG_SUBSCRIPTIONRUNTIMECRITICALEXCEPTION_OCCURRED_SUBSCRIPTIONCONSUMER_ARG_POLLING_TOPICS_ARG_C96324AD,
- this,
- topicNames,
- ex);
- lastSubscriptionRuntimeCriticalException = ex;
- } else {
- LOGGER.warn(
- SubscriptionMessages
-
.LOG_EXECUTIONEXCEPTION_OCCURRED_SUBSCRIPTIONCONSUMER_ARG_POLLING_TOPICS_ARG_40F5E1CC,
- this,
- topicNames,
- e);
- }
- }
- }
+ return collectMultiplePollResults(
+ SubscriptionExecutorServiceManager.submitMultiplePollTasks(tasks,
timeoutMs), topicNames);
} catch (final InterruptedException e) {
LOGGER.warn(
SubscriptionMessages
@@ -802,6 +814,56 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
// TODO: ignore possible interrupted state?
+ return Collections.emptyList();
+ }
+
+ List<SubscriptionMessage> collectMultiplePollResults(
+ final List<Future<List<SubscriptionMessage>>> futures, final Set<String>
topicNames)
+ throws InterruptedException {
+ final List<SubscriptionMessage> messages = new ArrayList<>();
+ SubscriptionRuntimeCriticalException
lastSubscriptionRuntimeCriticalException = null;
+ for (final Future<List<SubscriptionMessage>> future : futures) {
+ try {
+ if (future.isCancelled()) {
+ continue;
+ }
+ messages.addAll(future.get());
+ } catch (final CancellationException ignored) {
+
+ } catch (final ExecutionException e) {
+ final Throwable cause = e.getCause();
+ if (cause instanceof SubscriptionConsumerFencedException) {
+ final SubscriptionConsumerFencedException fencedException =
+ (SubscriptionConsumerFencedException) cause;
+ fence(fencedException);
+ throw fencedException;
+ }
+ if (cause instanceof SubscriptionRuntimeCriticalException) {
+ final SubscriptionRuntimeCriticalException ex =
+ (SubscriptionRuntimeCriticalException) cause;
+ LOGGER.warn(
+ SubscriptionMessages
+
.LOG_SUBSCRIPTIONRUNTIMECRITICALEXCEPTION_OCCURRED_SUBSCRIPTIONCONSUMER_ARG_POLLING_TOPICS_ARG_C96324AD,
+ this,
+ topicNames,
+ ex);
+ lastSubscriptionRuntimeCriticalException = ex;
+ } else {
+ LOGGER.warn(
+ SubscriptionMessages
+
.LOG_EXECUTIONEXCEPTION_OCCURRED_SUBSCRIPTIONCONSUMER_ARG_POLLING_TOPICS_ARG_40F5E1CC,
+ this,
+ topicNames,
+ e);
+ }
+ }
+ }
+
+ // A timed-out task can be cancelled after fencing the consumer but before
its exception is
+ // observable through Future#get. Never deliver messages collected by
sibling tasks in that
+ // case.
+ checkIfFenced();
+
// even if a SubscriptionRuntimeCriticalException is encountered, try to
deliver the message to
// the client
if (messages.isEmpty() &&
Objects.nonNull(lastSubscriptionRuntimeCriticalException)) {
@@ -830,6 +892,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
private List<SubscriptionMessage> singlePoll(
/* @NotNull */ final Set<String> topicNames, final long timeoutMs)
throws SubscriptionException {
+ checkIfFenced();
if (topicNames.isEmpty()) {
return Collections.emptyList();
}
@@ -876,6 +939,10 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
}
}
} catch (final SubscriptionRuntimeCriticalException e) {
+ if (e instanceof SubscriptionConsumerFencedException) {
+ fence((SubscriptionConsumerFencedException) e);
+ throw e;
+ }
LOGGER.warn(
SubscriptionMessages
.LOG_SUBSCRIPTIONRUNTIMECRITICALEXCEPTION_OCCURRED_SUBSCRIPTIONCONSUMER_ARG_POLLING_TOPICS_ARG_C96324AD,
@@ -974,6 +1041,10 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
try (final RandomAccessFile fileWriter = new RandomAccessFile(file, "rw"))
{
return pollFileInternal(commitContext, fileName, file, fileWriter,
timer);
} catch (final Exception e) {
+ if (e instanceof SubscriptionConsumerFencedException) {
+ fence((SubscriptionConsumerFencedException) e);
+ throw (SubscriptionConsumerFencedException) e;
+ }
if (!(e instanceof SubscriptionPollTimeoutException)) {
inFlightFilesCommitContextSet.remove(commitContext);
}
@@ -1184,6 +1255,10 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
try {
return pollTabletsInternal(response, timer);
} catch (final Exception e) {
+ if (e instanceof SubscriptionConsumerFencedException) {
+ fence((SubscriptionConsumerFencedException) e);
+ throw (SubscriptionConsumerFencedException) e;
+ }
// construct temporary message to nack
nack(
Collections.singletonList(
@@ -1323,6 +1398,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
private List<SubscriptionPollResponse> pollInternal(
final Set<String> topicNames, final long timeoutMs) throws
SubscriptionException {
+ checkIfFenced();
providers.acquireReadLock();
try {
final AbstractSubscriptionProvider provider =
providers.getNextAvailableProvider();
@@ -1428,6 +1504,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
private void commit(final Iterable<SubscriptionCommitContext>
commitContexts, final boolean nack)
throws SubscriptionException {
+ checkIfFenced();
final Map<Integer, List<SubscriptionCommitContext>>
dataNodeIdToSubscriptionCommitContexts =
new HashMap<>();
for (final SubscriptionCommitContext commitContext : commitContexts) {
@@ -1590,7 +1667,12 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
nack,
dataNodeId));
}
- return provider.commit(subscriptionCommitContexts, nack);
+ try {
+ return provider.commit(subscriptionCommitContexts, nack);
+ } catch (final SubscriptionConsumerFencedException e) {
+ fence(e);
+ throw e;
+ }
} finally {
providers.releaseReadLock();
}
@@ -1607,7 +1689,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
future[0] =
SubscriptionExecutorServiceManager.submitHeartbeatWorker(
() -> {
- if (isClosed()) {
+ if (isClosed() || isFenced()) {
if (Objects.nonNull(future[0])) {
future[0].cancel(false);
LOGGER.info(SubscriptionMessages.CONSUMER_CANCEL_HEARTBEAT_WORKER, this);
@@ -1627,7 +1709,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
future[0] =
SubscriptionExecutorServiceManager.submitEndpointsSyncer(
() -> {
- if (isClosed()) {
+ if (isClosed() || isFenced()) {
if (Objects.nonNull(future[0])) {
future[0].cancel(false);
LOGGER.info(SubscriptionMessages.CONSUMER_CANCEL_ENDPOINTS_SYNCER, this);
@@ -1661,11 +1743,11 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
@Override
public void run() {
- if (isClosed()) {
- return;
- }
-
try {
+ checkIfFenced();
+ if (isClosed()) {
+ return;
+ }
ack(messages);
callback.onComplete();
} catch (final Exception e) {
@@ -1678,11 +1760,11 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
final CompletableFuture<Void> future = new CompletableFuture<>();
SubscriptionExecutorServiceManager.submitAsyncCommitWorker(
() -> {
- if (isClosed()) {
- return;
- }
-
try {
+ checkIfFenced();
+ if (isClosed()) {
+ return;
+ }
ack(messages);
future.complete(null);
} catch (final Throwable e) {
@@ -1709,6 +1791,10 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
subscribedTopics = provider.subscribe(topicNames);
return;
} catch (final Exception e) {
+ if (e instanceof SubscriptionConsumerFencedException) {
+ fence((SubscriptionConsumerFencedException) e);
+ throw (SubscriptionConsumerFencedException) e;
+ }
if (e instanceof SubscriptionOwnerFencedException) {
throw (SubscriptionOwnerFencedException) e;
}
@@ -1750,6 +1836,10 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
subscribedTopics = provider.unsubscribe(topicNames);
return;
} catch (final Exception e) {
+ if (e instanceof SubscriptionConsumerFencedException) {
+ fence((SubscriptionConsumerFencedException) e);
+ throw (SubscriptionConsumerFencedException) e;
+ }
if (e instanceof SubscriptionPipeTimeoutException) {
// degrade exception to log for pipe timeout
LOGGER.warn(e.getMessage());
@@ -1795,6 +1885,10 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
try {
provider.seek(topicName, seekType, timestamp);
} catch (final Exception e) {
+ if (e instanceof SubscriptionConsumerFencedException) {
+ fence((SubscriptionConsumerFencedException) e);
+ throw (SubscriptionConsumerFencedException) e;
+ }
failedProviders.add(provider);
if (Objects.isNull(firstFailure)) {
firstFailure = e;
@@ -1836,6 +1930,10 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
try {
provider.seekToTopicProgress(topicName, topicProgress);
} catch (final Exception e) {
+ if (e instanceof SubscriptionConsumerFencedException) {
+ fence((SubscriptionConsumerFencedException) e);
+ throw (SubscriptionConsumerFencedException) e;
+ }
failedProviders.add(provider);
if (Objects.isNull(firstFailure)) {
firstFailure = e;
@@ -1878,6 +1976,10 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
try {
provider.seekAfterTopicProgress(topicName, topicProgress);
} catch (final Exception e) {
+ if (e instanceof SubscriptionConsumerFencedException) {
+ fence((SubscriptionConsumerFencedException) e);
+ throw (SubscriptionConsumerFencedException) e;
+ }
failedProviders.add(provider);
if (Objects.isNull(firstFailure)) {
firstFailure = e;
@@ -2148,6 +2250,10 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
try {
return provider.heartbeat().getEndPoints();
} catch (final Exception e) {
+ if (e instanceof SubscriptionConsumerFencedException) {
+ fence((SubscriptionConsumerFencedException) e);
+ throw (SubscriptionConsumerFencedException) e;
+ }
LOGGER.warn(
SubscriptionMessages
.LOG_ARG_FAILED_FETCH_ALL_ENDPOINTS_SUBSCRIPTION_PROVIDER_ARG_TRY_NEXT_25651CAD,
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java
index 92350097787..02a6de8a9bb 100644
---
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java
@@ -26,6 +26,7 @@ import
org.apache.iotdb.rpc.subscription.config.ConsumerConfig;
import org.apache.iotdb.rpc.subscription.config.ConsumerConstant;
import org.apache.iotdb.rpc.subscription.config.TopicConfig;
import
org.apache.iotdb.rpc.subscription.exception.SubscriptionConnectionException;
+import
org.apache.iotdb.rpc.subscription.exception.SubscriptionConsumerFencedException;
import org.apache.iotdb.rpc.subscription.exception.SubscriptionException;
import
org.apache.iotdb.rpc.subscription.exception.SubscriptionOwnerFencedException;
import
org.apache.iotdb.rpc.subscription.exception.SubscriptionPipeTimeoutException;
@@ -91,6 +92,7 @@ public abstract class AbstractSubscriptionProvider {
private String consumerId;
private String consumerGroupId;
+ private String consumerInstanceId;
private final String ownerId;
private final Long ownerEpoch;
@@ -189,6 +191,10 @@ public abstract class AbstractSubscriptionProvider {
return consumerGroupId;
}
+ void setConsumerInstanceId(final String consumerInstanceId) {
+ this.consumerInstanceId = consumerInstanceId;
+ }
+
TEndPoint getEndPoint() {
return endPoint;
}
@@ -206,6 +212,9 @@ public abstract class AbstractSubscriptionProvider {
final Map<String, String> consumerAttributes = new HashMap<>();
consumerAttributes.put(ConsumerConstant.CONSUMER_GROUP_ID_KEY,
consumerGroupId);
consumerAttributes.put(ConsumerConstant.CONSUMER_ID_KEY, consumerId);
+ if (consumerInstanceId != null) {
+ consumerAttributes.put(ConsumerConstant.CONSUMER_INSTANCE_ID_KEY,
consumerInstanceId);
+ }
if (ownerId != null) {
consumerAttributes.put(ConsumerConstant.OWNER_ID_KEY, ownerId);
}
@@ -327,6 +336,15 @@ public abstract class AbstractSubscriptionProvider {
}
}
+ synchronized void closeSession() throws IoTDBConnectionException {
+ try {
+ session.close();
+ } finally {
+ setUnavailable();
+ isClosed.set(true);
+ }
+ }
+
void closeInternal() throws SubscriptionException {
final TPipeSubscribeResp resp;
try {
@@ -711,6 +729,9 @@ public abstract class AbstractSubscriptionProvider {
LOGGER.warn(errorMessage);
throw new SubscriptionOwnerFencedException(errorMessage);
}
+ case 1919: // SUBSCRIPTION_CONSUMER_FENCED
+ LOGGER.warn(status.message);
+ throw new SubscriptionConsumerFencedException(status.message);
case 1900: // SUBSCRIPTION_VERSION_ERROR
case 1901: // SUBSCRIPTION_TYPE_ERROR
case 1909: // SUBSCRIPTION_MISSING_CONSUMER
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java
index 0f8e032bec2..3eaf6e6906e 100644
---
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java
@@ -22,6 +22,7 @@ package org.apache.iotdb.session.subscription.consumer.base;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.rpc.IoTDBConnectionException;
import
org.apache.iotdb.rpc.subscription.exception.SubscriptionConnectionException;
+import
org.apache.iotdb.rpc.subscription.exception.SubscriptionConsumerFencedException;
import org.apache.iotdb.rpc.subscription.exception.SubscriptionException;
import org.apache.iotdb.rpc.subscription.i18n.SubscriptionMessages;
import
org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext;
@@ -92,6 +93,9 @@ final class AbstractSubscriptionProviders {
try {
defaultProvider = consumer.constructProviderAndHandshake(endPoint);
+ } catch (final SubscriptionConsumerFencedException e) {
+ consumer.fence(e);
+ throw e;
} catch (final Exception e) {
connectionFailures.put(endPoint,
consumer.sanitizeConnectionFailureMessage(e));
connectionFailureCauses.add(e);
@@ -109,6 +113,9 @@ final class AbstractSubscriptionProviders {
final Map<Integer, TEndPoint> allEndPoints;
try {
allEndPoints = defaultProvider.heartbeat().getEndPoints();
+ } catch (final SubscriptionConsumerFencedException e) {
+ consumer.fence(e);
+ throw e;
} catch (final Exception e) {
LOGGER.warn(
SubscriptionMessages.LOG_ARG_FAILED_FETCH_ALL_ENDPOINTS_ARG_BECAUSE_ARG_2C9E11D4,
@@ -127,6 +134,9 @@ final class AbstractSubscriptionProviders {
final AbstractSubscriptionProvider provider;
try {
provider = consumer.constructProviderAndHandshake(entry.getValue());
+ } catch (final SubscriptionConsumerFencedException e) {
+ consumer.fence(e);
+ throw e;
} catch (final Exception e) {
LOGGER.warn(
SubscriptionMessages.LOG_ARG_FAILED_CREATE_CONNECTION_ARG_BECAUSE_ARG_E536E22A,
@@ -164,11 +174,20 @@ final class AbstractSubscriptionProviders {
/** Detaches and closes the current providers. Terminal consumer close may
call this lock-free. */
void closeProviders() {
+ closeProviders(true);
+ }
+
+ /** Detaches and closes the current providers. Terminal consumer close may
call this lock-free. */
+ void closeProviders(final boolean closeConsumer) {
final List<AbstractSubscriptionProvider> providers = getAllProviders();
subscriptionProviders.clear();
for (final AbstractSubscriptionProvider provider : providers) {
try {
- provider.close();
+ if (closeConsumer) {
+ provider.close();
+ } else {
+ provider.closeSession();
+ }
} catch (final Exception e) {
LOGGER.warn(SubscriptionMessages.PROVIDER_CLOSE_FAILED, provider, e,
e);
}
@@ -293,13 +312,13 @@ final class AbstractSubscriptionProviders {
/////////////////////////////// heartbeat ///////////////////////////////
void heartbeat(final AbstractSubscriptionConsumer consumer) {
- if (consumer.isClosed()) {
+ if (consumer.isClosed() || consumer.isFenced()) {
return;
}
acquireWriteLock();
try {
- if (consumer.isClosed()) {
+ if (consumer.isClosed() || consumer.isFenced()) {
return;
}
heartbeatInternal(consumer);
@@ -310,6 +329,9 @@ final class AbstractSubscriptionProviders {
private void heartbeatInternal(final AbstractSubscriptionConsumer consumer) {
for (final AbstractSubscriptionProvider provider : getAllProviders()) {
+ if (consumer.isFenced()) {
+ return;
+ }
try {
final List<SubscriptionCommitContext> processorBufferedCommitContexts =
consumer.getProcessorBufferedCommitContexts(provider.getDataNodeId());
@@ -326,6 +348,10 @@ final class AbstractSubscriptionProviders {
consumer.unsubscribe(topicName);
}
provider.setAvailable();
+ } catch (final SubscriptionConsumerFencedException e) {
+ consumer.fence(e);
+ provider.setUnavailable();
+ return;
} catch (final Exception e) {
LOGGER.warn(
SubscriptionMessages
@@ -342,13 +368,13 @@ final class AbstractSubscriptionProviders {
/////////////////////////////// sync endpoints
///////////////////////////////
void sync(final AbstractSubscriptionConsumer consumer) {
- if (consumer.isClosed()) {
+ if (consumer.isClosed() || consumer.isFenced()) {
return;
}
acquireWriteLock();
try {
- if (consumer.isClosed()) {
+ if (consumer.isClosed() || consumer.isFenced()) {
return;
}
syncInternal(consumer);
@@ -358,9 +384,15 @@ final class AbstractSubscriptionProviders {
}
private void syncInternal(final AbstractSubscriptionConsumer consumer) {
+ if (consumer.isFenced()) {
+ return;
+ }
if (hasNoAvailableProviders()) {
try {
openProviders(consumer);
+ } catch (final SubscriptionConsumerFencedException e) {
+ consumer.fence(e);
+ return;
} catch (final Exception e) {
LOGGER.warn(SubscriptionMessages.OPEN_PROVIDERS_FAILED, consumer, e,
e);
return;
@@ -370,6 +402,9 @@ final class AbstractSubscriptionProviders {
final Map<Integer, TEndPoint> allEndPoints;
try {
allEndPoints = consumer.fetchAllEndPointsWithRedirection();
+ } catch (final SubscriptionConsumerFencedException e) {
+ consumer.fence(e);
+ return;
} catch (final Exception e) {
LOGGER.warn(SubscriptionMessages.FETCH_ENDPOINTS_FAILED, consumer, e, e);
return;
@@ -377,6 +412,9 @@ final class AbstractSubscriptionProviders {
// add new providers or handshake existing providers
for (final Map.Entry<Integer, TEndPoint> entry : allEndPoints.entrySet()) {
+ if (consumer.isFenced()) {
+ return;
+ }
final AbstractSubscriptionProvider provider =
getProvider(entry.getKey());
if (Objects.isNull(provider)) {
// new provider
@@ -384,6 +422,9 @@ final class AbstractSubscriptionProviders {
final AbstractSubscriptionProvider newProvider;
try {
newProvider = consumer.constructProviderAndHandshake(endPoint);
+ } catch (final SubscriptionConsumerFencedException e) {
+ consumer.fence(e);
+ return;
} catch (final Exception e) {
LOGGER.warn(
SubscriptionMessages.LOG_ARG_FAILED_CREATE_CONNECTION_ARG_BECAUSE_ARG_E536E22A,
@@ -399,6 +440,10 @@ final class AbstractSubscriptionProviders {
try {
consumer.subscribedTopics = provider.heartbeat().getTopics();
provider.setAvailable();
+ } catch (final SubscriptionConsumerFencedException e) {
+ consumer.fence(e);
+ provider.setUnavailable();
+ return;
} catch (final Exception e) {
LOGGER.warn(
SubscriptionMessages
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPullConsumer.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPullConsumer.java
index 61d558cd829..b565d0d9293 100644
---
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPullConsumer.java
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPullConsumer.java
@@ -166,6 +166,13 @@ public abstract class AbstractSubscriptionPullConsumer
extends AbstractSubscript
return;
}
+ if (isFenced()) {
+ isClosed.set(true);
+ prepareClose();
+ super.close();
+ return;
+ }
+
List<SubscriptionMessage> drainedProcessorMessages =
Collections.emptyList();
if (!processors.isEmpty()) {
drainedProcessorMessages = drainProcessorPipeline();
@@ -640,7 +647,7 @@ public abstract class AbstractSubscriptionPullConsumer
extends AbstractSubscript
future[0] =
SubscriptionExecutorServiceManager.submitAutoCommitWorker(
() -> {
- if (isClosed()) {
+ if (isClosed() || isFenced()) {
if (Objects.nonNull(future[0])) {
future[0].cancel(false);
LOGGER.info(SubscriptionMessages.PULL_CONSUMER_CANCEL_AUTO_COMMIT, this);
@@ -656,7 +663,7 @@ public abstract class AbstractSubscriptionPullConsumer
extends AbstractSubscript
private class AutoCommitWorker implements Runnable {
@Override
public void run() {
- if (isClosed()) {
+ if (isClosed() || isFenced()) {
return;
}
@@ -690,6 +697,9 @@ public abstract class AbstractSubscriptionPullConsumer
extends AbstractSubscript
}
private void commitAllUncommittedMessages() {
+ if (isFenced()) {
+ return;
+ }
for (final Map.Entry<Long, Set<SubscriptionCommitContext>> entry :
uncommittedCommitContexts.entrySet()) {
try {
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPushConsumer.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPushConsumer.java
index a06a2e9cc4e..972b4262e4d 100644
---
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPushConsumer.java
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionPushConsumer.java
@@ -175,7 +175,7 @@ public abstract class AbstractSubscriptionPushConsumer
extends AbstractSubscript
future[0] =
SubscriptionExecutorServiceManager.submitAutoPollWorker(
() -> {
- if (isClosed()) {
+ if (isClosed() || isFenced()) {
if (Objects.nonNull(future[0])) {
future[0].cancel(false);
LOGGER.info(SubscriptionMessages.PUSH_CONSUMER_CANCEL_AUTO_POLL, this);
@@ -191,7 +191,7 @@ public abstract class AbstractSubscriptionPushConsumer
extends AbstractSubscript
class AutoPollWorker implements Runnable {
@Override
public void run() {
- if (isClosed()) {
+ if (isClosed() || isFenced()) {
return;
}
diff --git
a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerLifecycleTest.java
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerLifecycleTest.java
index df42e2b244c..00db8d31145 100644
---
a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerLifecycleTest.java
+++
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerLifecycleTest.java
@@ -20,23 +20,39 @@
package org.apache.iotdb.session.subscription.consumer.base;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import
org.apache.iotdb.rpc.subscription.exception.SubscriptionConsumerFencedException;
import org.apache.iotdb.rpc.subscription.exception.SubscriptionException;
import
org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext;
+import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollResponse;
+import
org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollResponseType;
+import org.apache.iotdb.rpc.subscription.payload.poll.TabletsPayload;
+import org.apache.iotdb.rpc.subscription.payload.poll.TopicProgress;
import
org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeHeartbeatResp;
import org.apache.iotdb.session.AbstractSessionBuilder;
import org.apache.iotdb.session.subscription.SubscriptionTreeSessionBuilder;
+import org.apache.iotdb.session.subscription.consumer.AsyncCommitCallback;
+import org.apache.iotdb.session.subscription.payload.SubscriptionMessage;
import org.junit.Assert;
import org.junit.Test;
+import java.lang.reflect.Field;
import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
import java.util.List;
+import java.util.Map;
import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
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;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BooleanSupplier;
public class SubscriptionConsumerLifecycleTest {
@@ -97,6 +113,179 @@ public class SubscriptionConsumerLifecycleTest {
Assert.assertTrue(consumer.closedStatesDuringClose.get(0));
}
+ @Test
+ public void testFencedHeartbeatStopsBackgroundReconnect() throws Exception {
+ final TestPullConsumer consumer = new TestPullConsumer();
+ final AbstractSubscriptionProviders providers = getProviders(consumer);
+ try {
+ consumer.open();
+ consumer.fenceOnHeartbeat = true;
+ providers.heartbeat(consumer);
+
+ Assert.assertTrue(consumer.isFenced());
+ providers.sync(consumer);
+ providers.heartbeat(consumer);
+ Assert.assertEquals(1, consumer.createdProviders.size());
+ try {
+ consumer.multiplePoll(Collections.singleton("topic"), 100L);
+ Assert.fail("The fenced consumer must not poll or reconnect");
+ } catch (final SubscriptionConsumerFencedException expected) {
+ Assert.assertEquals("consumer connection fenced",
expected.getMessage());
+ }
+ consumer.close();
+ Assert.assertEquals(0, consumer.closeRequestCount);
+ Assert.assertEquals(1, consumer.sessionCloseCount);
+ Assert.assertEquals(1, consumer.closedStatesDuringClose.size());
+ } finally {
+ consumer.close();
+ }
+ }
+
+ @Test
+ public void testFencedDuringOpenClosesPartiallyOpenedProviders() throws
Exception {
+ final TestPullConsumer consumer = new TestPullConsumer();
+ consumer.fenceOnHeartbeat = true;
+ try {
+ consumer.open();
+ Assert.fail("The consumer must fail to open when its handshake is
fenced");
+ } catch (final SubscriptionConsumerFencedException expected) {
+ Assert.assertTrue(consumer.isFenced());
+ Assert.assertEquals(1, consumer.createdProviders.size());
+ Assert.assertEquals(0, consumer.closeRequestCount);
+ Assert.assertEquals(1, consumer.sessionCloseCount);
+ Assert.assertEquals(1, consumer.closedStatesDuringClose.size());
+ }
+
+ try {
+ consumer.open();
+ Assert.fail("The fenced consumer must not retry the handshake");
+ } catch (final SubscriptionConsumerFencedException expected) {
+ Assert.assertEquals(1, consumer.createdProviders.size());
+ }
+ }
+
+ @Test
+ public void testFencedHandshakeClosesOpenedSession() throws Exception {
+ final TestPullConsumer consumer = new TestPullConsumer();
+ consumer.fenceOnHandshake = true;
+
+ try {
+ consumer.open();
+ Assert.fail("The consumer must fail to open when its handshake is
fenced");
+ } catch (final SubscriptionConsumerFencedException expected) {
+ Assert.assertTrue(consumer.isFenced());
+ Assert.assertEquals(1, consumer.createdProviders.size());
+ Assert.assertEquals(0, consumer.closeRequestCount);
+ Assert.assertEquals(1, consumer.sessionCloseCount);
+ Assert.assertEquals(1, consumer.closedStatesDuringClose.size());
+ }
+ }
+
+ @Test
+ public void testFencedTabletContinuationDoesNotSendNack() throws Exception {
+ final TestPullConsumer consumer = new TestPullConsumer();
+ try {
+ consumer.open();
+ consumer.returnPartialTablets = true;
+ consumer.fenceOnPollTablets = true;
+
+ try {
+ consumer.multiplePoll(Collections.singleton("topic"), 1_000L);
+ Assert.fail("A fenced tablet continuation must fail the poll");
+ } catch (final SubscriptionConsumerFencedException expected) {
+ Assert.assertEquals("consumer connection fenced",
expected.getMessage());
+ }
+
+ Assert.assertTrue(consumer.isFenced());
+ Assert.assertEquals(0, consumer.commitRequestCount);
+ } finally {
+ consumer.close();
+ }
+ }
+
+ @Test
+ public void testFencedParallelPollDoesNotDeliverSiblingMessages() throws
Exception {
+ final TestPullConsumer consumer = new TestPullConsumer();
+ final SubscriptionConsumerFencedException fencedException =
+ new SubscriptionConsumerFencedException("consumer connection fenced");
+ final SubscriptionMessage message =
+ new SubscriptionMessage(
+ new SubscriptionCommitContext(0, 0, "topic", CONSUMER_GROUP_ID,
0L), 1L);
+ final CompletableFuture<List<SubscriptionMessage>> fencedFuture = new
CompletableFuture<>();
+ fencedFuture.completeExceptionally(fencedException);
+
+ try {
+ consumer.collectMultiplePollResults(
+ Arrays.asList(
+
CompletableFuture.completedFuture(Collections.singletonList(message)),
fencedFuture),
+ Collections.singleton("topic"));
+ Assert.fail("A fenced poll task must discard messages returned by
sibling tasks");
+ } catch (final SubscriptionConsumerFencedException expected) {
+ Assert.assertSame(fencedException, expected);
+ }
+
+ Assert.assertTrue(consumer.isFenced());
+ }
+
+ @Test
+ public void testFencedAsyncCommitFailsBeforeReadingMessages() throws
Exception {
+ final TestPullConsumer consumer = new TestPullConsumer();
+ try {
+ final SubscriptionConsumerFencedException fencedException =
+ new SubscriptionConsumerFencedException("consumer connection
fenced");
+ consumer.fence(fencedException);
+ final AtomicBoolean messagesIterated = new AtomicBoolean(false);
+ final Iterable<SubscriptionMessage> messages =
+ () -> {
+ messagesIterated.set(true);
+ return Collections.emptyIterator();
+ };
+
+ final CountDownLatch callbackCompleted = new CountDownLatch(1);
+ final AtomicBoolean callbackSucceeded = new AtomicBoolean(false);
+ final AtomicReference<Throwable> callbackFailure = new
AtomicReference<>();
+ consumer.commitAsync(
+ messages,
+ new AsyncCommitCallback() {
+ @Override
+ public void onComplete() {
+ callbackSucceeded.set(true);
+ callbackCompleted.countDown();
+ }
+
+ @Override
+ public void onFailure(final Throwable e) {
+ callbackFailure.set(e);
+ callbackCompleted.countDown();
+ }
+ });
+
+ Assert.assertTrue(callbackCompleted.await(5, TimeUnit.SECONDS));
+ Assert.assertFalse(callbackSucceeded.get());
+ Assert.assertSame(fencedException, callbackFailure.get());
+
+ final CompletableFuture<Void> future = consumer.commitAsync(messages);
+ try {
+ future.get(5, TimeUnit.SECONDS);
+ Assert.fail("A fenced async commit must complete exceptionally");
+ } catch (final ExecutionException expected) {
+ Assert.assertSame(fencedException, expected.getCause());
+ }
+
+ Assert.assertFalse(messagesIterated.get());
+ Assert.assertEquals(0, consumer.commitRequestCount);
+ } finally {
+ consumer.close();
+ }
+ }
+
+ private AbstractSubscriptionProviders getProviders(final
AbstractSubscriptionConsumer consumer)
+ throws Exception {
+ final Field field =
AbstractSubscriptionConsumer.class.getDeclaredField("providers");
+ field.setAccessible(true);
+ return (AbstractSubscriptionProviders) field.get(consumer);
+ }
+
@Test
public void testConcurrentPullConsumerCloseReturnsWithoutWaiting() throws
Exception {
final CountDownLatch providerCloseStarted = new CountDownLatch(1);
@@ -169,7 +358,16 @@ public class SubscriptionConsumerLifecycleTest {
connectionTimeoutInMs,
this::isClosed,
closedStatesDuringHandshake,
- closedStatesDuringClose);
+ closedStatesDuringClose,
+ () -> false,
+ () -> false,
+ () -> false,
+ () -> false,
+ () -> {},
+ () -> {},
+ () -> {},
+ null,
+ null);
}
}
@@ -177,6 +375,14 @@ public class SubscriptionConsumerLifecycleTest {
private final List<Boolean> closedStatesDuringHandshake = new
ArrayList<>();
private final List<Boolean> closedStatesDuringClose = new ArrayList<>();
+ private final List<TestSubscriptionProvider> createdProviders = new
ArrayList<>();
+ private boolean fenceOnHandshake;
+ private boolean fenceOnHeartbeat;
+ private boolean fenceOnPollTablets;
+ private boolean returnPartialTablets;
+ private int closeRequestCount;
+ private int sessionCloseCount;
+ private int commitRequestCount;
private final CountDownLatch providerCloseStarted;
private final CountDownLatch allowProviderClose;
@@ -212,23 +418,33 @@ public class SubscriptionConsumerLifecycleTest {
final int thriftMaxFrameSize,
final long heartbeatIntervalMs,
final int connectionTimeoutInMs) {
- return new TestSubscriptionProvider(
- endPoint,
- username,
- password,
- encryptedPassword,
- consumerId,
- consumerGroupId,
- ownerId,
- ownerEpoch,
- thriftMaxFrameSize,
- heartbeatIntervalMs,
- connectionTimeoutInMs,
- this::isClosed,
- closedStatesDuringHandshake,
- closedStatesDuringClose,
- providerCloseStarted,
- allowProviderClose);
+ final TestSubscriptionProvider provider =
+ new TestSubscriptionProvider(
+ endPoint,
+ username,
+ password,
+ encryptedPassword,
+ consumerId,
+ consumerGroupId,
+ ownerId,
+ ownerEpoch,
+ thriftMaxFrameSize,
+ heartbeatIntervalMs,
+ connectionTimeoutInMs,
+ this::isClosed,
+ closedStatesDuringHandshake,
+ closedStatesDuringClose,
+ () -> fenceOnHandshake,
+ () -> fenceOnHeartbeat,
+ () -> fenceOnPollTablets,
+ () -> returnPartialTablets,
+ () -> closeRequestCount++,
+ () -> commitRequestCount++,
+ () -> sessionCloseCount++,
+ providerCloseStarted,
+ allowProviderClose);
+ createdProviders.add(provider);
+ return provider;
}
}
@@ -237,43 +453,16 @@ public class SubscriptionConsumerLifecycleTest {
private final BooleanSupplier consumerClosedSupplier;
private final List<Boolean> closedStatesDuringHandshake;
private final List<Boolean> closedStatesDuringClose;
+ private final BooleanSupplier fenceOnHandshake;
+ private final BooleanSupplier fenceOnHeartbeat;
+ private final BooleanSupplier fenceOnPollTablets;
+ private final BooleanSupplier returnPartialTablets;
+ private final Runnable closeRequest;
+ private final Runnable commitRequest;
+ private final Runnable sessionClose;
private final CountDownLatch providerCloseStarted;
private final CountDownLatch allowProviderClose;
- private TestSubscriptionProvider(
- final TEndPoint endPoint,
- final String username,
- final String password,
- final String encryptedPassword,
- final String consumerId,
- final String consumerGroupId,
- final String ownerId,
- final Long ownerEpoch,
- final int thriftMaxFrameSize,
- final long heartbeatIntervalMs,
- final int connectionTimeoutInMs,
- final BooleanSupplier consumerClosedSupplier,
- final List<Boolean> closedStatesDuringHandshake,
- final List<Boolean> closedStatesDuringClose) {
- this(
- endPoint,
- username,
- password,
- encryptedPassword,
- consumerId,
- consumerGroupId,
- ownerId,
- ownerEpoch,
- thriftMaxFrameSize,
- heartbeatIntervalMs,
- connectionTimeoutInMs,
- consumerClosedSupplier,
- closedStatesDuringHandshake,
- closedStatesDuringClose,
- null,
- null);
- }
-
private TestSubscriptionProvider(
final TEndPoint endPoint,
final String username,
@@ -289,6 +478,13 @@ public class SubscriptionConsumerLifecycleTest {
final BooleanSupplier consumerClosedSupplier,
final List<Boolean> closedStatesDuringHandshake,
final List<Boolean> closedStatesDuringClose,
+ final BooleanSupplier fenceOnHandshake,
+ final BooleanSupplier fenceOnHeartbeat,
+ final BooleanSupplier fenceOnPollTablets,
+ final BooleanSupplier returnPartialTablets,
+ final Runnable closeRequest,
+ final Runnable commitRequest,
+ final Runnable sessionClose,
final CountDownLatch providerCloseStarted,
final CountDownLatch allowProviderClose) {
super(
@@ -306,6 +502,13 @@ public class SubscriptionConsumerLifecycleTest {
this.consumerClosedSupplier = consumerClosedSupplier;
this.closedStatesDuringHandshake = closedStatesDuringHandshake;
this.closedStatesDuringClose = closedStatesDuringClose;
+ this.fenceOnHandshake = fenceOnHandshake;
+ this.fenceOnHeartbeat = fenceOnHeartbeat;
+ this.fenceOnPollTablets = fenceOnPollTablets;
+ this.returnPartialTablets = returnPartialTablets;
+ this.closeRequest = closeRequest;
+ this.commitRequest = commitRequest;
+ this.sessionClose = sessionClose;
this.providerCloseStarted = providerCloseStarted;
this.allowProviderClose = allowProviderClose;
}
@@ -333,12 +536,16 @@ public class SubscriptionConsumerLifecycleTest {
@Override
synchronized void handshake() {
closedStatesDuringHandshake.add(consumerClosedSupplier.getAsBoolean());
+ if (fenceOnHandshake.getAsBoolean()) {
+ throw new SubscriptionConsumerFencedException("consumer connection
fenced");
+ }
setAvailable();
}
@Override
synchronized void close() {
closedStatesDuringClose.add(consumerClosedSupplier.getAsBoolean());
+ closeRequest.run();
if (Objects.nonNull(providerCloseStarted)) {
providerCloseStarted.countDown();
}
@@ -352,10 +559,51 @@ public class SubscriptionConsumerLifecycleTest {
setUnavailable();
}
+ @Override
+ synchronized void closeSession() {
+ closedStatesDuringClose.add(consumerClosedSupplier.getAsBoolean());
+ sessionClose.run();
+ setUnavailable();
+ }
+
@Override
PipeSubscribeHeartbeatResp heartbeat(
final List<SubscriptionCommitContext> processorBufferedCommitContexts)
{
+ if (fenceOnHeartbeat.getAsBoolean()) {
+ throw new SubscriptionConsumerFencedException("consumer connection
fenced");
+ }
return new PipeSubscribeHeartbeatResp();
}
+
+ @Override
+ List<SubscriptionPollResponse> poll(
+ final Set<String> topicNames,
+ final long timeoutMs,
+ final Map<String, TopicProgress> progressByTopic) {
+ if (!returnPartialTablets.getAsBoolean()) {
+ return Collections.emptyList();
+ }
+ return Collections.singletonList(
+ new SubscriptionPollResponse(
+ SubscriptionPollResponseType.TABLETS.getType(),
+ new TabletsPayload(Collections.emptyMap(), 1),
+ new SubscriptionCommitContext(0, 0, "topic", CONSUMER_GROUP_ID,
0L)));
+ }
+
+ @Override
+ List<SubscriptionPollResponse> pollTablets(
+ final SubscriptionCommitContext commitContext, final int offset, final
long timeoutMs) {
+ if (fenceOnPollTablets.getAsBoolean()) {
+ throw new SubscriptionConsumerFencedException("consumer connection
fenced");
+ }
+ return Collections.emptyList();
+ }
+
+ @Override
+ CommitResult commit(
+ final List<SubscriptionCommitContext> subscriptionCommitContexts,
final boolean nack) {
+ commitRequest.run();
+ return CommitResult.empty();
+ }
}
}
diff --git
a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionProviderStatusTest.java
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionProviderStatusTest.java
new file mode 100644
index 00000000000..5da15cadcb4
--- /dev/null
+++
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionProviderStatusTest.java
@@ -0,0 +1,72 @@
+/*
+ * 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.session.subscription.consumer.base;
+
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.rpc.TSStatusCode;
+import
org.apache.iotdb.rpc.subscription.exception.SubscriptionConsumerFencedException;
+import org.apache.iotdb.rpc.subscription.exception.SubscriptionException;
+import
org.apache.iotdb.rpc.subscription.exception.SubscriptionRuntimeCriticalException;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+
+public class SubscriptionProviderStatusTest {
+
+ @Test
+ public void testConsumerFencedStatusMapsToSpecificException() throws
Exception {
+ final SubscriptionException exception =
+ invokeVerifyPipeSubscribeSuccess(
+ new
TSStatus(TSStatusCode.SUBSCRIPTION_CONSUMER_FENCED.getStatusCode())
+ .setMessage("consumer fenced"));
+
+ Assert.assertTrue(exception instanceof
SubscriptionConsumerFencedException);
+ Assert.assertEquals("consumer fenced", exception.getMessage());
+ }
+
+ @Test
+ public void testMissingConsumerStatusRemainsCriticalException() throws
Exception {
+ final SubscriptionException exception =
+ invokeVerifyPipeSubscribeSuccess(
+ new
TSStatus(TSStatusCode.SUBSCRIPTION_MISSING_CONSUMER.getStatusCode())
+ .setMessage("missing consumer"));
+
+ Assert.assertEquals(SubscriptionRuntimeCriticalException.class,
exception.getClass());
+ Assert.assertEquals("missing consumer", exception.getMessage());
+ }
+
+ private SubscriptionException invokeVerifyPipeSubscribeSuccess(final
TSStatus status)
+ throws Exception {
+ final Method method =
+ AbstractSubscriptionProvider.class.getDeclaredMethod(
+ "verifyPipeSubscribeSuccess", TSStatus.class);
+ method.setAccessible(true);
+ try {
+ method.invoke(null, status);
+ Assert.fail("Expected a subscription exception");
+ return null;
+ } catch (final InvocationTargetException e) {
+ return (SubscriptionException) e.getCause();
+ }
+ }
+}
diff --git
a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
index b818a52e838..d0236c1293d 100644
---
a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
+++
b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
@@ -1798,6 +1798,12 @@ public final class DataNodePipeMessages {
"Subscription: The consumer {} has already existed when handshaking,
skip creating consumer.";
public static final String
PIPE_LOG_SUBSCRIPTION_CONSUMER_HANDSHAKE_SUCCESSFULLY_DATA_NODE_ID_58DA6A5F =
"Subscription: consumer {} handshake successfully, data node id: {}";
+ public static final String
+
LOG_SUBSCRIPTION_CONSUMER_ARG_IN_CONSUMER_GROUP_ARG_WAS_TAKEN_OVER_BY_A_NEWER_CONNECTION_FENCED_THE_PREVIOUS_CONNECTION_4E72DBD9
=
+ "Subscription: consumer {} in consumer group {} was taken over by a
newer connection; fenced the previous connection.";
+ public static final String
+
MESSAGE_SUBSCRIPTION_CONSUMER_CONNECTION_WAS_FENCED_BECAUSE_A_NEWER_CONNECTION_WITH_THE_SAME_CONSUMER_ID_AND_CONSUMER_GROUP_ID_COMPLETED_THE_HANDSHAKE_THIS_CONSUMER_INSTANCE_CANNOT_BE_REUSED_CREATE_A_NEW_CONSUMER_INSTANCE_TO_RECONNECT_B0C2CCBE
=
+ "Subscription: consumer connection was fenced because a newer
connection with the same consumer ID and consumer group ID completed the
handshake. This consumer instance cannot be reused; create a new consumer
instance to reconnect.";
public static final String
PIPE_LOG_SUBSCRIPTION_CONSUMER_UNSUBSCRIBE_SUCCESSFULLY_AA5E0AA9 =
"Subscription: consumer {} unsubscribe {} successfully";
public static final String
PIPE_LOG_SUBSCRIPTION_CONSUMER_COMMIT_NACK_ACCEPTED_SUCCESSFULLY_58D1C111 =
diff --git
a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
index fe768a1f018..d11e16f41f6 100644
---
a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
+++
b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
@@ -1667,6 +1667,12 @@ public final class DataNodePipeMessages {
"Subscription：握手时 consumer {} 已存在，跳过 consumer 创建。";
public static final String
PIPE_LOG_SUBSCRIPTION_CONSUMER_HANDSHAKE_SUCCESSFULLY_DATA_NODE_ID_58DA6A5F =
"Subscription：consumer {} 握手成功，data node id：{}";
+ public static final String
+
LOG_SUBSCRIPTION_CONSUMER_ARG_IN_CONSUMER_GROUP_ARG_WAS_TAKEN_OVER_BY_A_NEWER_CONNECTION_FENCED_THE_PREVIOUS_CONNECTION_4E72DBD9
=
+ "Subscription：consumer {}（consumer group {}）已被新连接接管，旧连接已被隔离。";
+ public static final String
+
MESSAGE_SUBSCRIPTION_CONSUMER_CONNECTION_WAS_FENCED_BECAUSE_A_NEWER_CONNECTION_WITH_THE_SAME_CONSUMER_ID_AND_CONSUMER_GROUP_ID_COMPLETED_THE_HANDSHAKE_THIS_CONSUMER_INSTANCE_CANNOT_BE_REUSED_CREATE_A_NEW_CONSUMER_INSTANCE_TO_RECONNECT_B0C2CCBE
=
+ "Subscription：consumer 连接已被隔离，因为具有相同 consumer ID 和 consumer group ID
的新连接已完成握手。当前 consumer 实例不可复用，请创建新的 consumer 实例进行重连。";
public static final String
PIPE_LOG_SUBSCRIPTION_CONSUMER_UNSUBSCRIBE_SUCCESSFULLY_AA5E0AA9 =
"Subscription：consumer {} 取消订阅 {} 成功";
public static final String
PIPE_LOG_SUBSCRIPTION_CONSUMER_COMMIT_NACK_ACCEPTED_SUCCESSFULLY_58D1C111 =
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 3fdadb6e84e..dc42d4fbcae 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
@@ -47,6 +47,7 @@ import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.BooleanSupplier;
import java.util.function.Supplier;
@@ -62,6 +63,15 @@ public class SubscriptionReceiverAgent {
PipeSubscribeResponseVersion.VERSION_1.getVersion(),
PipeSubscribeResponseType.ACK.getType());
+ private static final TPipeSubscribeResp SUBSCRIPTION_CONSUMER_FENCED_RESP =
+ new TPipeSubscribeResp(
+ RpcUtils.getStatus(
+ TSStatusCode.SUBSCRIPTION_CONSUMER_FENCED,
+ DataNodePipeMessages
+
.MESSAGE_SUBSCRIPTION_CONSUMER_CONNECTION_WAS_FENCED_BECAUSE_A_NEWER_CONNECTION_WITH_THE_SAME_CONSUMER_ID_AND_CONSUMER_GROUP_ID_COMPLETED_THE_HANDSHAKE_THIS_CONSUMER_INSTANCE_CANNOT_BE_REUSED_CREATE_A_NEW_CONSUMER_INSTANCE_TO_RECONNECT_B0C2CCBE),
+ PipeSubscribeResponseVersion.VERSION_1.getVersion(),
+ PipeSubscribeResponseType.ACK.getType());
+
private final Map<Byte, Supplier<SubscriptionReceiver>> receiverConstructors
= new HashMap<>();
private final ThreadLocal<SubscriptionReceiver> receiverThreadLocal = new
ThreadLocal<>();
@@ -127,7 +137,8 @@ public class SubscriptionReceiverAgent {
if (receiverConstructors.containsKey(reqVersion)) {
final SubscriptionReceiver receiver = getReceiver(reqVersion);
receiver.setAuthenticatedUsername(username);
- final ConsumerIdentity consumerIdentity = getConsumerIdentity(req,
receiver);
+ final ConsumerConnection consumerConnection = getConsumerConnection(req,
receiver);
+ final ConsumerIdentity consumerIdentity = consumerConnection.identity();
final RequestResult requestResult = new RequestResult();
if (Objects.isNull(consumerIdentity)) {
@@ -136,12 +147,21 @@ public class SubscriptionReceiverAgent {
consumerReceivers.compute(
consumerIdentity,
(identity, currentReceiver) -> {
+ if (isHandshake(req)
+ && currentReceiver != null
+ && currentReceiver != receiver
+ && shouldKeepCurrentReceiver(
+ currentReceiver,
consumerConnection.consumerInstanceId())) {
+ receiver.invalidateConsumer();
+ requestResult.response = SUBSCRIPTION_CONSUMER_FENCED_RESP;
+ return currentReceiver;
+ }
requestResult.response = handleRequest(receiver, req,
currentReceiver);
if (isHandshake(req)) {
if (isSuccessful(requestResult.response)) {
if (currentReceiver != null && currentReceiver != receiver) {
- currentReceiver.invalidateConsumer();
+ invalidateReplacedReceiver(currentReceiver, identity);
}
return receiver;
}
@@ -158,7 +178,9 @@ public class SubscriptionReceiverAgent {
if (isHandshake(req) && isSuccessful(requestResult.response)) {
final ConsumerIdentity activeIdentity = getConsumerIdentity(receiver);
if (!Objects.equals(consumerIdentity, activeIdentity)) {
- registerReceiver(receiver, activeIdentity);
+ if (!registerReceiver(receiver, activeIdentity)) {
+ requestResult.response = SUBSCRIPTION_CONSUMER_FENCED_RESP;
+ }
} else {
removeReceiverMappingsExcept(receiver, activeIdentity);
}
@@ -281,21 +303,36 @@ public class SubscriptionReceiverAgent {
return receiver.handle(req);
}
- private void registerReceiver(
+ private boolean registerReceiver(
final SubscriptionReceiver receiver, final ConsumerIdentity identity) {
if (Objects.isNull(identity)) {
removeReceiverMappings(receiver);
- return;
+ return true;
}
+ final AtomicBoolean registered = new AtomicBoolean(false);
consumerReceivers.compute(
identity,
(key, currentReceiver) -> {
- if (currentReceiver != null && currentReceiver != receiver) {
- currentReceiver.invalidateConsumer();
+ if (currentReceiver == null || currentReceiver == receiver) {
+ registered.set(true);
+ return receiver;
+ }
+ if (receiver.getConsumerInstanceId() != null
+ && !shouldKeepCurrentReceiver(currentReceiver,
receiver.getConsumerInstanceId())) {
+ invalidateReplacedReceiver(currentReceiver, key);
+ registered.set(true);
+ return receiver;
}
- return receiver;
+ // The receiver completed its handshake after another receiver had
already claimed the
+ // identity. Keep the current owner and fence this late receiver
instead of allowing an
+ // old connection to take the consumer back.
+ invalidateReplacedReceiver(receiver, key);
+ return currentReceiver;
});
- removeReceiverMappingsExcept(receiver, identity);
+ if (registered.get()) {
+ removeReceiverMappingsExcept(receiver, identity);
+ }
+ return registered.get();
}
private void removeReceiverMappingsExcept(
@@ -314,7 +351,17 @@ public class SubscriptionReceiverAgent {
(identity, currentReceiver) -> consumerReceivers.remove(identity,
receiver));
}
- private static ConsumerIdentity getConsumerIdentity(
+ private void invalidateReplacedReceiver(
+ final SubscriptionReceiver receiver, final ConsumerIdentity identity) {
+ LOGGER.info(
+ DataNodePipeMessages
+
.LOG_SUBSCRIPTION_CONSUMER_ARG_IN_CONSUMER_GROUP_ARG_WAS_TAKEN_OVER_BY_A_NEWER_CONNECTION_FENCED_THE_PREVIOUS_CONNECTION_4E72DBD9,
+ identity.consumerId(),
+ identity.consumerGroupId());
+ receiver.invalidateConsumer();
+ }
+
+ private static ConsumerConnection getConsumerConnection(
final TPipeSubscribeReq req, final SubscriptionReceiver receiver) {
if (isHandshake(req) && req.isSetBody()) {
try {
@@ -325,7 +372,7 @@ public class SubscriptionReceiverAgent {
ConsumerIdentity.of(
consumerConfig.getConsumerGroupId(),
consumerConfig.getConsumerId());
if (Objects.nonNull(identity)) {
- return identity;
+ return new ConsumerConnection(identity,
consumerConfig.getConsumerInstanceId());
}
}
} catch (final RuntimeException ignored) {
@@ -333,7 +380,20 @@ public class SubscriptionReceiverAgent {
// original buffer, so parsing is intentionally done on a duplicate
above.
}
}
- return getConsumerIdentity(receiver);
+ return new ConsumerConnection(getConsumerIdentity(receiver),
receiver.getConsumerInstanceId());
+ }
+
+ private static boolean shouldKeepCurrentReceiver(
+ final SubscriptionReceiver currentReceiver, final String
incomingConsumerInstanceId) {
+ final String currentConsumerInstanceId =
currentReceiver.getConsumerInstanceId();
+ if (Objects.equals(currentConsumerInstanceId, incomingConsumerInstanceId))
{
+ return false;
+ }
+ if (currentConsumerInstanceId == null) {
+ return false;
+ }
+ return incomingConsumerInstanceId == null
+ || currentConsumerInstanceId.compareTo(incomingConsumerInstanceId) > 0;
}
private static ConsumerIdentity getConsumerIdentity(final
SubscriptionReceiver receiver) {
@@ -365,4 +425,6 @@ public class SubscriptionReceiverAgent {
: new ConsumerIdentity(consumerGroupId, consumerId);
}
}
+
+ private record ConsumerConnection(ConsumerIdentity identity, String
consumerInstanceId) {}
}
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 e5de617c89c..1d2fbe4fc18 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
@@ -47,6 +47,14 @@ public interface SubscriptionReceiver {
*/
String getConsumerGroupId();
+ /**
+ * Returns the identifier shared by all DataNode connections of the current
consumer instance, or
+ * {@code null} for a legacy client.
+ */
+ default String getConsumerInstanceId() {
+ return null;
+ }
+
/**
* Invalidates this receiver so that requests from an obsolete connection
cannot affect a new
* owner.
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 b0eec505a5c..4c8c1c08d30 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
@@ -47,6 +47,7 @@ import
org.apache.iotdb.db.subscription.metric.SubscriptionPrefetchingQueueMetri
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.config.TopicConfig;
import org.apache.iotdb.rpc.subscription.exception.SubscriptionException;
import
org.apache.iotdb.rpc.subscription.exception.SubscriptionPayloadExceedException;
@@ -121,11 +122,21 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
PipeSubscribeResponseVersion.VERSION_1.getVersion(),
PipeSubscribeResponseType.ACK.getType());
+ private static final TPipeSubscribeResp SUBSCRIPTION_CONSUMER_FENCED_RESP =
+ new TPipeSubscribeResp(
+ RpcUtils.getStatus(
+ TSStatusCode.SUBSCRIPTION_CONSUMER_FENCED,
+ DataNodePipeMessages
+
.MESSAGE_SUBSCRIPTION_CONSUMER_CONNECTION_WAS_FENCED_BECAUSE_A_NEWER_CONNECTION_WITH_THE_SAME_CONSUMER_ID_AND_CONSUMER_GROUP_ID_COMPLETED_THE_HANDSHAKE_THIS_CONSUMER_INSTANCE_CANNOT_BE_REUSED_CREATE_A_NEW_CONSUMER_INSTANCE_TO_RECONNECT_B0C2CCBE),
+ PipeSubscribeResponseVersion.VERSION_1.getVersion(),
+ PipeSubscribeResponseType.ACK.getType());
+
private final ThreadLocal<ConsumerConfig> consumerConfigThreadLocal = new
ThreadLocal<>();
private final ThreadLocal<PollTimer> pollTimerThreadLocal = new
ThreadLocal<>();
private volatile String authenticatedUsername;
private volatile ConsumerConfig sharedConsumerConfig;
private volatile boolean consumerInvalidated;
+ private volatile boolean consumerFenced;
private volatile long lastActivityTimeMs = System.currentTimeMillis();
private final AtomicLong inFlightRequestCount = new AtomicLong(0);
private long consumerStateVersion;
@@ -135,8 +146,11 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
@Override
public final TPipeSubscribeResp handle(final TPipeSubscribeReq req) {
final short reqType = req.getType();
- beforeHandle(reqType);
+ final boolean isFencedRequest = beforeHandle(reqType);
try {
+ if (isFencedRequest) {
+ return SUBSCRIPTION_CONSUMER_FENCED_RESP;
+ }
if (PipeSubscribeRequestType.isValidatedRequestType(reqType)) {
switch (PipeSubscribeRequestType.valueOf(reqType)) {
case HANDSHAKE:
@@ -199,6 +213,12 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
return Objects.isNull(consumerConfig) ? null :
consumerConfig.getConsumerGroupId();
}
+ @Override
+ public String getConsumerInstanceId() {
+ final ConsumerConfig consumerConfig = sharedConsumerConfig;
+ return Objects.isNull(consumerConfig) ? null :
consumerConfig.getConsumerInstanceId();
+ }
+
@Override
public boolean hasActiveConsumer() {
return Objects.nonNull(sharedConsumerConfig);
@@ -206,7 +226,10 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
@Override
public void invalidateConsumer() {
- clearSharedConsumerState();
+ synchronized (this) {
+ consumerFenced = true;
+ clearSharedConsumerState();
+ }
}
@Override
@@ -1192,11 +1215,14 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
private void createConsumer(final ConsumerConfig consumerConfig) throws
SubscriptionException {
try (final ConfigNodeClient configNodeClient =
CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
+ final Map<String, String> persistedConsumerAttributes =
+ new HashMap<>(consumerConfig.getAttribute());
+
persistedConsumerAttributes.remove(ConsumerConstant.CONSUMER_INSTANCE_ID_KEY);
final TCreateConsumerReq req =
new TCreateConsumerReq()
.setConsumerId(consumerConfig.getConsumerId())
.setConsumerGroupId(consumerConfig.getConsumerGroupId())
- .setConsumerAttributes(consumerConfig.getAttribute());
+ .setConsumerAttributes(persistedConsumerAttributes);
final TSStatus tsStatus = configNodeClient.createConsumer(req);
if (TSStatusCode.SUCCESS_STATUS.getStatusCode() != tsStatus.getCode()) {
LOGGER.warn(
@@ -1340,17 +1366,23 @@ public class SubscriptionReceiverV1 implements
SubscriptionReceiver {
}
}
- private void beforeHandle(final short reqType) {
+ private boolean beforeHandle(final short reqType) {
synchronized (this) {
+ final boolean isHandshake = PipeSubscribeRequestType.HANDSHAKE.getType()
== reqType;
+ // A receiver fenced by a newer connection is terminal. In particular,
do not allow the old
+ // connection to handshake again and reclaim the consumer identity. A
normal disconnected
+ // receiver remains recoverable through the existing consumerInvalidated
handshake path.
+ final boolean isFencedRequest = consumerFenced;
if (consumerInvalidated) {
consumerConfigThreadLocal.remove();
pollTimerThreadLocal.remove();
- if (PipeSubscribeRequestType.HANDSHAKE.getType() == reqType) {
+ if (isHandshake && !consumerFenced) {
consumerInvalidated = false;
}
}
inFlightRequestCount.incrementAndGet();
lastActivityTimeMs = System.currentTimeMillis();
+ return isFencedRequest;
}
}
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 bdf75a3a339..2d48518e5f2 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
@@ -29,6 +29,8 @@ 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.request.PipeSubscribeSubscribeReq;
+import
org.apache.iotdb.rpc.subscription.payload.request.SubscriptionHeartbeatReq;
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;
@@ -41,6 +43,7 @@ import java.io.IOException;
import java.lang.reflect.Field;
import java.util.HashMap;
import java.util.Map;
+import java.util.Set;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ScheduledExecutorService;
@@ -135,6 +138,291 @@ public class SubscriptionReceiverAgentTest {
Assert.assertEquals(1, newReceiver.timeoutCount.get());
}
+ @Test
+ public void
testDuplicateConnectionFencesOldReceiverWithoutInvalidatingNewReceiver()
+ throws Exception {
+ final CopyOnWriteArrayList<FakeSubscriptionReceiver> receivers = new
CopyOnWriteArrayList<>();
+ final SubscriptionReceiverAgent agent = createAgent(receivers, false /*
closeOnTimeout */);
+ final TPipeSubscribeReq handshake = createHandshakeRequest("group",
"consumer");
+
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ agent.handle(handshake, "root").getStatus().getCode());
+ final AtomicReference<TPipeSubscribeResp> duplicateResponse = new
AtomicReference<>();
+ final Thread newConnection =
+ new Thread(() -> duplicateResponse.set(agent.handle(handshake,
"root")));
+ newConnection.start();
+ newConnection.join(TimeUnit.SECONDS.toMillis(10));
+ Assert.assertFalse(newConnection.isAlive());
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
duplicateResponse.get().getStatus().getCode());
+ final FakeSubscriptionReceiver oldReceiver = receivers.get(0);
+ final FakeSubscriptionReceiver newReceiver = receivers.get(1);
+ Assert.assertTrue(oldReceiver.invalidated);
+ Assert.assertNotNull(newReceiver.consumerConfig);
+ Assert.assertEquals(
+ TSStatusCode.SUBSCRIPTION_CONSUMER_FENCED.getStatusCode(),
+ agent.handle(SubscriptionHeartbeatReq.toThriftReq(),
"root").getStatus().getCode());
+
+ agent.checkReceiverTimeouts();
+
+ Assert.assertEquals(0, oldReceiver.timeoutCount.get());
+ Assert.assertEquals(1, newReceiver.timeoutCount.get());
+ }
+
+ @Test
+ public void testConsumerInstanceWinnerIsIndependentOfHandshakeOrder() throws
Exception {
+ final TPipeSubscribeReq olderHandshake =
+ createHandshakeRequest("group", "consumer", "0000000000000001-older");
+ final TPipeSubscribeReq newerHandshake =
+ createHandshakeRequest("group", "consumer", "0000000000000002-newer");
+
+ final CopyOnWriteArrayList<FakeSubscriptionReceiver> olderFirstReceivers =
+ new CopyOnWriteArrayList<>();
+ final SubscriptionReceiverAgent olderFirstAgent =
+ createAgent(olderFirstReceivers, false /* closeOnTimeout */);
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ handleOnNewConnection(olderFirstAgent,
olderHandshake).getStatus().getCode());
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ handleOnNewConnection(olderFirstAgent,
newerHandshake).getStatus().getCode());
+ Assert.assertTrue(olderFirstReceivers.get(0).invalidated);
+ Assert.assertFalse(olderFirstReceivers.get(1).invalidated);
+
+ final CopyOnWriteArrayList<FakeSubscriptionReceiver> newerFirstReceivers =
+ new CopyOnWriteArrayList<>();
+ final SubscriptionReceiverAgent newerFirstAgent =
+ createAgent(newerFirstReceivers, false /* closeOnTimeout */);
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ handleOnNewConnection(newerFirstAgent,
newerHandshake).getStatus().getCode());
+ Assert.assertEquals(
+ TSStatusCode.SUBSCRIPTION_CONSUMER_FENCED.getStatusCode(),
+ handleOnNewConnection(newerFirstAgent,
olderHandshake).getStatus().getCode());
+ Assert.assertFalse(newerFirstReceivers.get(0).invalidated);
+ Assert.assertTrue(newerFirstReceivers.get(1).invalidated);
+ }
+
+ @Test
+ public void testConcurrentHandshakeWithSameIdentityFencesOldReceiver()
throws Exception {
+ final CopyOnWriteArrayList<FakeSubscriptionReceiver> receivers = new
CopyOnWriteArrayList<>();
+ final CountDownLatch oldHandshakeEntered = new CountDownLatch(1);
+ final CountDownLatch releaseOldHandshake = new CountDownLatch(1);
+ final CountDownLatch newReceiverCreated = new CountDownLatch(1);
+ final CountDownLatch newHandshakeFinished = new CountDownLatch(1);
+ final CountDownLatch oldHeartbeatFinished = new CountDownLatch(1);
+ final AtomicInteger receiverIndex = new AtomicInteger();
+ final AtomicReference<Throwable> threadFailure = new AtomicReference<>();
+ final SubscriptionReceiverAgent agent =
+ new SubscriptionReceiverAgent(
+ () -> {
+ final boolean isOldReceiver = receiverIndex.getAndIncrement() ==
0;
+ final FakeSubscriptionReceiver receiver =
+ new FakeSubscriptionReceiver(
+ false,
+ false,
+ isOldReceiver ? oldHandshakeEntered : null,
+ isOldReceiver ? releaseOldHandshake : null);
+ receivers.add(receiver);
+ if (!isOldReceiver) {
+ newReceiverCreated.countDown();
+ }
+ return receiver;
+ },
+ false,
+ () -> true);
+ final TPipeSubscribeReq oldHandshake = createHandshakeRequest("group",
"consumer");
+ final TPipeSubscribeReq newHandshake = createHandshakeRequest("group",
"consumer");
+ final AtomicReference<TPipeSubscribeResp> oldHandshakeResponse = new
AtomicReference<>();
+ final AtomicReference<TPipeSubscribeResp> newHandshakeResponse = new
AtomicReference<>();
+ final AtomicReference<TPipeSubscribeResp> oldHeartbeatResponse = new
AtomicReference<>();
+ final AtomicReference<TPipeSubscribeResp> newHeartbeatResponse = new
AtomicReference<>();
+ final AtomicReference<TPipeSubscribeResp> newSubscribeResponse = new
AtomicReference<>();
+
+ final Thread oldConnection =
+ new Thread(
+ () -> {
+ try {
+ oldHandshakeResponse.set(agent.handle(oldHandshake, "root"));
+ if (!newHandshakeFinished.await(10, TimeUnit.SECONDS)) {
+ throw new AssertionError("The new handshake did not finish");
+ }
+ oldHeartbeatResponse.set(
+ agent.handle(SubscriptionHeartbeatReq.toThriftReq(),
"root"));
+ } catch (final Throwable t) {
+ threadFailure.compareAndSet(null, t);
+ } finally {
+ oldHeartbeatFinished.countDown();
+ }
+ });
+ final Thread newConnection =
+ new Thread(
+ () -> {
+ try {
+ newHandshakeResponse.set(agent.handle(newHandshake, "root"));
+ newHandshakeFinished.countDown();
+ if (!oldHeartbeatFinished.await(10, TimeUnit.SECONDS)) {
+ throw new AssertionError("The old heartbeat did not finish");
+ }
+ newHeartbeatResponse.set(
+ agent.handle(SubscriptionHeartbeatReq.toThriftReq(),
"root"));
+ newSubscribeResponse.set(
+ agent.handle(
+
PipeSubscribeSubscribeReq.toTPipeSubscribeReq(Set.of("topic")), "root"));
+ } catch (final Throwable t) {
+ threadFailure.compareAndSet(null, t);
+ } finally {
+ newHandshakeFinished.countDown();
+ }
+ });
+
+ oldConnection.start();
+ Assert.assertTrue(oldHandshakeEntered.await(10, TimeUnit.SECONDS));
+ newConnection.start();
+ try {
+ Assert.assertTrue(newReceiverCreated.await(10, TimeUnit.SECONDS));
+ } finally {
+ releaseOldHandshake.countDown();
+ oldConnection.join(TimeUnit.SECONDS.toMillis(10));
+ newConnection.join(TimeUnit.SECONDS.toMillis(10));
+ }
+
+ Assert.assertFalse(oldConnection.isAlive());
+ Assert.assertFalse(newConnection.isAlive());
+ if (threadFailure.get() != null) {
+ throw new AssertionError(threadFailure.get());
+ }
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ oldHandshakeResponse.get().getStatus().getCode());
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ newHandshakeResponse.get().getStatus().getCode());
+ Assert.assertEquals(
+ TSStatusCode.SUBSCRIPTION_CONSUMER_FENCED.getStatusCode(),
+ oldHeartbeatResponse.get().getStatus().getCode());
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ newHeartbeatResponse.get().getStatus().getCode());
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ newSubscribeResponse.get().getStatus().getCode());
+ Assert.assertEquals(2, receivers.size());
+ Assert.assertTrue(receivers.get(0).invalidated);
+ Assert.assertFalse(receivers.get(1).invalidated);
+ }
+
+ @Test
+ public void testReconnectSucceedsAfterActiveConnectionExits() throws
IOException {
+ final CopyOnWriteArrayList<FakeSubscriptionReceiver> receivers = new
CopyOnWriteArrayList<>();
+ final SubscriptionReceiverAgent agent = createAgent(receivers, false /*
closeOnTimeout */);
+ final TPipeSubscribeReq handshake = createHandshakeRequest("group",
"consumer");
+
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ agent.handle(handshake, "root").getStatus().getCode());
+ agent.handleClientExit();
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ agent.handle(handshake, "root").getStatus().getCode());
+
+ Assert.assertTrue(receivers.get(0).invalidated);
+ Assert.assertFalse(receivers.get(1).invalidated);
+ }
+
+ @Test
+ public void testLateHandshakeCannotTakeOverNewReceiver() throws Exception {
+ final CopyOnWriteArrayList<FakeSubscriptionReceiver> receivers = new
CopyOnWriteArrayList<>();
+ final CountDownLatch oldHandshakeEntered = new CountDownLatch(1);
+ final CountDownLatch releaseOldHandshake = new CountDownLatch(1);
+ final CountDownLatch newOwnerReady = new CountDownLatch(1);
+ final CountDownLatch oldHandshakeFinished = new CountDownLatch(1);
+ final AtomicInteger receiverIndex = new AtomicInteger();
+ final AtomicReference<Throwable> threadFailure = new AtomicReference<>();
+ final SubscriptionReceiverAgent agent =
+ new SubscriptionReceiverAgent(
+ () -> {
+ final boolean isOldReceiver = receiverIndex.getAndIncrement() ==
0;
+ final FakeSubscriptionReceiver receiver =
+ new FakeSubscriptionReceiver(
+ false,
+ true,
+ isOldReceiver ? oldHandshakeEntered : null,
+ isOldReceiver ? releaseOldHandshake : null);
+ receivers.add(receiver);
+ return receiver;
+ },
+ false,
+ () -> true);
+ final TPipeSubscribeReq handshake =
createHandshakeRequestWithoutIdentity();
+ final AtomicReference<TPipeSubscribeResp> oldHandshakeResponse = new
AtomicReference<>();
+ final AtomicReference<TPipeSubscribeResp> newHandshakeResponse = new
AtomicReference<>();
+ final AtomicReference<TPipeSubscribeResp> newHeartbeatResponse = new
AtomicReference<>();
+ final AtomicReference<TPipeSubscribeResp> newSubscribeResponse = new
AtomicReference<>();
+
+ final Thread oldConnection =
+ new Thread(
+ () -> {
+ try {
+ oldHandshakeResponse.set(agent.handle(handshake, "root"));
+ } catch (final Throwable t) {
+ threadFailure.set(t);
+ } finally {
+ oldHandshakeFinished.countDown();
+ }
+ });
+ oldConnection.start();
+ Assert.assertTrue(oldHandshakeEntered.await(10, TimeUnit.SECONDS));
+
+ final Thread newConnection =
+ new Thread(
+ () -> {
+ try {
+ newHandshakeResponse.set(agent.handle(handshake, "root"));
+ newOwnerReady.countDown();
+ oldHandshakeFinished.await(10, TimeUnit.SECONDS);
+ newHeartbeatResponse.set(
+ agent.handle(SubscriptionHeartbeatReq.toThriftReq(),
"root"));
+ newSubscribeResponse.set(
+ agent.handle(
+
PipeSubscribeSubscribeReq.toTPipeSubscribeReq(Set.of("topic")), "root"));
+ } catch (final Throwable t) {
+ threadFailure.set(t);
+ newOwnerReady.countDown();
+ }
+ });
+ newConnection.start();
+ try {
+ Assert.assertTrue(newOwnerReady.await(10, TimeUnit.SECONDS));
+ } finally {
+ releaseOldHandshake.countDown();
+ oldConnection.join(TimeUnit.SECONDS.toMillis(10));
+ newConnection.join(TimeUnit.SECONDS.toMillis(10));
+ }
+
+ Assert.assertFalse(oldConnection.isAlive());
+ Assert.assertFalse(newConnection.isAlive());
+ if (threadFailure.get() != null) {
+ throw new AssertionError(threadFailure.get());
+ }
+ Assert.assertEquals(
+ TSStatusCode.SUBSCRIPTION_CONSUMER_FENCED.getStatusCode(),
+ oldHandshakeResponse.get().getStatus().getCode());
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ newHandshakeResponse.get().getStatus().getCode());
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ newHeartbeatResponse.get().getStatus().getCode());
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ newSubscribeResponse.get().getStatus().getCode());
+ Assert.assertEquals(2, receivers.size());
+ Assert.assertTrue(receivers.get(0).invalidated);
+ Assert.assertFalse(receivers.get(1).invalidated);
+ }
+
@Test
public void testLateExitFromOldConnectionKeepsNewReceiverRegistered() throws
Exception {
final CopyOnWriteArrayList<FakeSubscriptionReceiver> receivers = new
CopyOnWriteArrayList<>();
@@ -228,28 +516,90 @@ public class SubscriptionReceiverAgentTest {
private TPipeSubscribeReq createHandshakeRequest(
final String consumerGroupId, final String consumerId) throws
IOException {
+ return createHandshakeRequest(consumerGroupId, consumerId, null);
+ }
+
+ private TPipeSubscribeReq createHandshakeRequest(
+ final String consumerGroupId, final String consumerId, final String
consumerInstanceId)
+ throws IOException {
final Map<String, String> attributes = new HashMap<>();
attributes.put(ConsumerConstant.CONSUMER_GROUP_ID_KEY, consumerGroupId);
attributes.put(ConsumerConstant.CONSUMER_ID_KEY, consumerId);
+ if (consumerInstanceId != null) {
+ attributes.put(ConsumerConstant.CONSUMER_INSTANCE_ID_KEY,
consumerInstanceId);
+ }
return PipeSubscribeHandshakeReq.toTPipeSubscribeReq(new
ConsumerConfig(attributes));
}
+ private TPipeSubscribeResp handleOnNewConnection(
+ final SubscriptionReceiverAgent agent, final TPipeSubscribeReq request)
throws Exception {
+ final AtomicReference<TPipeSubscribeResp> response = new
AtomicReference<>();
+ final AtomicReference<Throwable> failure = new AtomicReference<>();
+ final Thread connection =
+ new Thread(
+ () -> {
+ try {
+ response.set(agent.handle(request, "root"));
+ } catch (final Throwable t) {
+ failure.set(t);
+ }
+ });
+ connection.start();
+ connection.join(TimeUnit.SECONDS.toMillis(10));
+ Assert.assertFalse(connection.isAlive());
+ if (failure.get() != null) {
+ throw new AssertionError(failure.get());
+ }
+ return response.get();
+ }
+
+ private TPipeSubscribeReq createHandshakeRequestWithoutIdentity() throws
IOException {
+ return PipeSubscribeHandshakeReq.toTPipeSubscribeReq(new
ConsumerConfig(new HashMap<>()));
+ }
+
private static class FakeSubscriptionReceiver implements
SubscriptionReceiver {
private final boolean closeOnTimeout;
+ private final boolean assignDefaultIdentity;
+ private final CountDownLatch handshakeEntered;
+ private final CountDownLatch releaseHandshake;
private final AtomicInteger timeoutCount = new AtomicInteger();
private final AtomicInteger exitCount = new AtomicInteger();
- private ConsumerConfig consumerConfig;
- private boolean invalidated;
+ private volatile ConsumerConfig consumerConfig;
+ private volatile boolean invalidated;
private FakeSubscriptionReceiver(final boolean closeOnTimeout) {
+ this(closeOnTimeout, false, null, null);
+ }
+
+ private FakeSubscriptionReceiver(
+ final boolean closeOnTimeout,
+ final boolean assignDefaultIdentity,
+ final CountDownLatch handshakeEntered,
+ final CountDownLatch releaseHandshake) {
this.closeOnTimeout = closeOnTimeout;
+ this.assignDefaultIdentity = assignDefaultIdentity;
+ this.handshakeEntered = handshakeEntered;
+ this.releaseHandshake = releaseHandshake;
}
@Override
public TPipeSubscribeResp handle(final TPipeSubscribeReq req) {
if (req.getType() == PipeSubscribeRequestType.HANDSHAKE.getType()) {
consumerConfig = ConsumerConfig.deserialize(req.bufferForBody());
+ if (assignDefaultIdentity) {
+ consumerConfig.setConsumerGroupId("group");
+ consumerConfig.setConsumerId("consumer");
+ }
+ if (handshakeEntered != null) {
+ handshakeEntered.countDown();
+ try {
+ Assert.assertTrue(releaseHandshake.await(10, TimeUnit.SECONDS));
+ } catch (final InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new AssertionError(e);
+ }
+ }
invalidated = false;
return response(TSStatusCode.SUCCESS_STATUS);
}
@@ -259,7 +609,7 @@ public class SubscriptionReceiverAgentTest {
return response(TSStatusCode.SUCCESS_STATUS);
}
return response(
- invalidated ? TSStatusCode.SUBSCRIPTION_MISSING_CONSUMER :
TSStatusCode.SUCCESS_STATUS);
+ invalidated ? TSStatusCode.SUBSCRIPTION_CONSUMER_FENCED :
TSStatusCode.SUCCESS_STATUS);
}
@Override
@@ -296,6 +646,11 @@ public class SubscriptionReceiverAgentTest {
return consumerConfig == null ? null :
consumerConfig.getConsumerGroupId();
}
+ @Override
+ public String getConsumerInstanceId() {
+ return consumerConfig == null ? null :
consumerConfig.getConsumerInstanceId();
+ }
+
@Override
public void invalidateConsumer() {
consumerConfig = null;
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 21d2e7b67d8..807d7231514 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
@@ -26,6 +26,8 @@ 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.config.TopicConstant;
+import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeHandshakeReq;
+import
org.apache.iotdb.rpc.subscription.payload.request.SubscriptionHeartbeatReq;
import org.junit.Assert;
import org.junit.Test;
@@ -159,6 +161,45 @@ public class SubscriptionReceiverV1Test {
Assert.assertNull(consumerConfigThreadLocal.get());
}
+ @Test
+ public void testInvalidatedConsumerReturnsFencedStatus() throws Exception {
+ final SubscriptionReceiverV1 receiver = new SubscriptionReceiverV1();
+ final ConsumerConfig consumerConfig = createConsumerConfig(1_000L);
+ setField(receiver, "sharedConsumerConfig", consumerConfig);
+
+ receiver.invalidateConsumer();
+
+ Assert.assertEquals(
+ TSStatusCode.SUBSCRIPTION_CONSUMER_FENCED.getStatusCode(),
+
receiver.handle(SubscriptionHeartbeatReq.toThriftReq()).getStatus().getCode());
+ }
+
+ @Test
+ public void testFencedConsumerCannotHandshakeAgain() throws Exception {
+ final SubscriptionReceiverV1 receiver = new SubscriptionReceiverV1();
+ final ConsumerConfig consumerConfig = createConsumerConfig(1_000L);
+
+ setField(receiver, "sharedConsumerConfig", consumerConfig);
+ receiver.invalidateConsumer();
+
+ Assert.assertEquals(
+ TSStatusCode.SUBSCRIPTION_CONSUMER_FENCED.getStatusCode(),
+ receiver
+
.handle(PipeSubscribeHandshakeReq.toTPipeSubscribeReq(consumerConfig))
+ .getStatus()
+ .getCode());
+ Assert.assertTrue((boolean) getField(receiver, "consumerFenced"));
+ }
+
+ @Test
+ public void testNeverHandshakenConsumerStillReturnsMissingConsumerStatus() {
+ final SubscriptionReceiverV1 receiver = new SubscriptionReceiverV1();
+
+ Assert.assertEquals(
+ TSStatusCode.SUBSCRIPTION_MISSING_CONSUMER.getStatusCode(),
+
receiver.handle(SubscriptionHeartbeatReq.toThriftReq()).getStatus().getCode());
+ }
+
@Test
public void testCalculateConsumerInactivityTimeoutUsesDefaultTimeout()
throws Exception {
final SubscriptionReceiverV1 receiver = new SubscriptionReceiverV1();
