From 703668d15f5bca8de688c2842296964133e3adb6 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Wed, 16 Sep 2026 19:27:39 +0800 Subject: [PATCH 1/2] Fix concurrent subscription consumer handshakes --- .../org/apache/iotdb/rpc/TSStatusCode.java | 1 + .../SubscriptionConsumerFencedException.java | 51 +++ .../base/AbstractSubscriptionConsumer.java | 83 ++++- .../base/AbstractSubscriptionProvider.java | 13 + .../base/AbstractSubscriptionProviders.java | 55 +++- .../AbstractSubscriptionPullConsumer.java | 13 +- .../AbstractSubscriptionPushConsumer.java | 4 +- .../SubscriptionConsumerLifecycleTest.java | 231 +++++++++++++- .../base/SubscriptionProviderStatusTest.java | 72 +++++ .../iotdb/db/i18n/DataNodePipeMessages.java | 6 + .../iotdb/db/i18n/DataNodePipeMessages.java | 6 + .../agent/SubscriptionReceiverAgent.java | 47 ++- .../receiver/SubscriptionReceiverV1.java | 30 +- .../agent/SubscriptionReceiverAgentTest.java | 291 +++++++++++++++++- .../receiver/SubscriptionReceiverV1Test.java | 41 +++ 15 files changed, 898 insertions(+), 46 deletions(-) create mode 100644 iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/exception/SubscriptionConsumerFencedException.java create mode 100644 iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionProviderStatusTest.java 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 61431249eb371..fce965fe27be4 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/exception/SubscriptionConsumerFencedException.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/exception/SubscriptionConsumerFencedException.java new file mode 100644 index 0000000000000..ef8e83b2c7c1c --- /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 73a51aa9923d5..15aaf6e43d324 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.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,7 @@ import java.util.concurrent.Future; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import java.util.stream.Collectors; @@ -123,6 +125,8 @@ abstract class AbstractSubscriptionConsumer implements AutoCloseable { private final AbstractSubscriptionProviders providers; private final AtomicBoolean isClosed = new AtomicBoolean(true); + private final AtomicReference fencedException = + new AtomicReference<>(); // This variable indicates whether the consumer has ever been closed. private final AtomicBoolean isReleased = new AtomicBoolean(false); @@ -296,6 +300,7 @@ protected AbstractSubscriptionConsumer( /////////////////////////////// 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 +310,7 @@ private void checkIfHasBeenClosed() throws SubscriptionException { } 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 +330,9 @@ protected synchronized void open() throws SubscriptionException { providers.acquireWriteLock(); try { providers.openProviders(this); // throw SubscriptionException + } catch (final SubscriptionException e) { + providers.closeProviders(!isFenced()); + throw e; } finally { providers.releaseWriteLock(); } @@ -347,7 +356,7 @@ public synchronized void close() { // close subscription providers providers.acquireWriteLock(); try { - providers.closeProviders(); + providers.closeProviders(!isFenced()); } finally { providers.releaseWriteLock(); } @@ -362,6 +371,21 @@ 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 { @@ -562,9 +586,13 @@ AbstractSubscriptionProvider constructProviderAndHandshake(final TEndPoint endPo 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 @@ -735,6 +763,7 @@ public long getLatestWatermarkTimestamp() { protected List multiplePoll( /* @NotNull */ final Set topicNames, final long timeoutMs) { + checkIfFenced(); if (topicNames.isEmpty()) { return Collections.emptyList(); } @@ -831,6 +860,7 @@ public List call() { private List singlePoll( /* @NotNull */ final Set topicNames, final long timeoutMs) throws SubscriptionException { + checkIfFenced(); if (topicNames.isEmpty()) { return Collections.emptyList(); } @@ -877,6 +907,10 @@ private List singlePoll( } } } 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, @@ -975,6 +1009,10 @@ private Optional pollFile( 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); } @@ -1185,6 +1223,10 @@ private Optional pollTablets( 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( @@ -1324,6 +1366,7 @@ private static void mergeTimeSelectedByTable( private List pollInternal( final Set topicNames, final long timeoutMs) throws SubscriptionException { + checkIfFenced(); providers.acquireReadLock(); try { final AbstractSubscriptionProvider provider = providers.getNextAvailableProvider(); @@ -1429,6 +1472,7 @@ private Iterable extractCommitContexts( private void commit(final Iterable commitContexts, final boolean nack) throws SubscriptionException { + checkIfFenced(); final Map> dataNodeIdToSubscriptionCommitContexts = new HashMap<>(); for (final SubscriptionCommitContext commitContext : commitContexts) { @@ -1591,7 +1635,12 @@ private AbstractSubscriptionProvider.CommitResult commitInternal( nack, dataNodeId)); } - return provider.commit(subscriptionCommitContexts, nack); + try { + return provider.commit(subscriptionCommitContexts, nack); + } catch (final SubscriptionConsumerFencedException e) { + fence(e); + throw e; + } } finally { providers.releaseReadLock(); } @@ -1608,7 +1657,7 @@ private void submitHeartbeatWorker() { 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); @@ -1628,7 +1677,7 @@ private void submitEndpointsSyncer() { 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); @@ -1710,6 +1759,10 @@ private void subscribeWithRedirection(final Set topicNames) throws Subsc 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; } @@ -1751,6 +1804,10 @@ private void unsubscribeWithRedirection(final Set topicNames) 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()); @@ -1796,6 +1853,10 @@ private void seekOnAllProviders( 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; @@ -1837,6 +1898,10 @@ private void seekToTopicProgressOnAllProviders( 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; @@ -1879,6 +1944,10 @@ private void seekAfterTopicProgressOnAllProviders( 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; @@ -2149,6 +2218,10 @@ Map fetchAllEndPointsWithRedirection() throws SubscriptionEx 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 cfe741e2d7e21..2316f79aa95d4 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.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; @@ -263,6 +264,15 @@ synchronized void close() throws SubscriptionException, IoTDBConnectionException } } + synchronized void closeSession() throws IoTDBConnectionException { + try { + session.close(); + } finally { + setUnavailable(); + isClosed.set(true); + } + } + void closeInternal() throws SubscriptionException { final TPipeSubscribeResp resp; try { @@ -647,6 +657,9 @@ private static void verifyPipeSubscribeSuccess(final TSStatus status) 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 d89fe998cdf2c..70be70eddb9cc 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 @@ 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; @@ -90,6 +91,9 @@ void openProviders(final AbstractSubscriptionConsumer consumer) throws Subscript 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); @@ -107,6 +111,9 @@ void openProviders(final AbstractSubscriptionConsumer consumer) throws Subscript final Map 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, @@ -125,6 +132,9 @@ void openProviders(final AbstractSubscriptionConsumer consumer) throws Subscript 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, @@ -162,9 +172,18 @@ void openProviders(final AbstractSubscriptionConsumer consumer) throws Subscript /** Caller should ensure that the method is called in the lock {@link #acquireWriteLock()}. */ void closeProviders() { + closeProviders(true); + } + + /** Caller should ensure that the method is called in the lock {@link #acquireWriteLock()}. */ + void closeProviders(final boolean closeConsumer) { for (final AbstractSubscriptionProvider provider : getAllProviders()) { try { - provider.close(); + if (closeConsumer) { + provider.close(); + } else { + provider.closeSession(); + } } catch (final Exception e) { LOGGER.warn(SubscriptionMessages.PROVIDER_CLOSE_FAILED, provider, e, e); } @@ -263,13 +282,13 @@ AbstractSubscriptionProvider getNextAvailableProvider() { /////////////////////////////// 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); @@ -280,6 +299,9 @@ void heartbeat(final AbstractSubscriptionConsumer consumer) { private void heartbeatInternal(final AbstractSubscriptionConsumer consumer) { for (final AbstractSubscriptionProvider provider : getAllProviders()) { + if (consumer.isFenced()) { + return; + } try { final List processorBufferedCommitContexts = consumer.getProcessorBufferedCommitContexts(provider.getDataNodeId()); @@ -296,6 +318,10 @@ private void heartbeatInternal(final AbstractSubscriptionConsumer consumer) { consumer.unsubscribe(topicName); } provider.setAvailable(); + } catch (final SubscriptionConsumerFencedException e) { + consumer.fence(e); + provider.setUnavailable(); + return; } catch (final Exception e) { LOGGER.warn( SubscriptionMessages @@ -312,13 +338,13 @@ private void heartbeatInternal(final AbstractSubscriptionConsumer consumer) { /////////////////////////////// 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); @@ -328,9 +354,15 @@ void sync(final AbstractSubscriptionConsumer consumer) { } 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; @@ -340,6 +372,9 @@ private void syncInternal(final AbstractSubscriptionConsumer consumer) { final Map 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; @@ -347,6 +382,9 @@ private void syncInternal(final AbstractSubscriptionConsumer consumer) { // add new providers or handshake existing providers for (final Map.Entry entry : allEndPoints.entrySet()) { + if (consumer.isFenced()) { + return; + } final AbstractSubscriptionProvider provider = getProvider(entry.getKey()); if (Objects.isNull(provider)) { // new provider @@ -354,6 +392,9 @@ private void syncInternal(final AbstractSubscriptionConsumer consumer) { 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, @@ -369,6 +410,10 @@ private void syncInternal(final AbstractSubscriptionConsumer consumer) { 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 811b4228c1e8a..7c0313e021b00 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 @@ -159,6 +159,12 @@ public synchronized void close() { return; } + if (isFenced()) { + isClosed.set(true); + super.close(); + return; + } + if (!processors.isEmpty()) { if (autoCommit) { final List drainedMessages = drainProcessorPipeline(); @@ -626,7 +632,7 @@ private void submitAutoCommitWorker() { 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); @@ -642,7 +648,7 @@ private void submitAutoCommitWorker() { private class AutoCommitWorker implements Runnable { @Override public void run() { - if (isClosed()) { + if (isClosed() || isFenced()) { return; } @@ -676,6 +682,9 @@ public void run() { } private void commitAllUncommittedMessages() { + if (isFenced()) { + return; + } for (final Map.Entry> 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 f3c7d38c4dc7b..9ed60564fdbf6 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 @@ -163,7 +163,7 @@ private void submitAutoPollWorker() { 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); @@ -179,7 +179,7 @@ private void submitAutoPollWorker() { 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 68f714f2ce52c..abc01f9875172 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,8 +20,13 @@ 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; @@ -29,9 +34,13 @@ import org.junit.Assert; import org.junit.Test; +import java.lang.reflect.Field; import java.util.ArrayList; +import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.Objects; +import java.util.Set; import java.util.function.BooleanSupplier; public class SubscriptionConsumerLifecycleTest { @@ -92,6 +101,103 @@ public void testPullConsumerIsClosedBeforeProviderClose() throws SubscriptionExc 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(); + } + } + + private AbstractSubscriptionProviders getProviders(final AbstractSubscriptionConsumer consumer) + throws Exception { + final Field field = AbstractSubscriptionConsumer.class.getDeclaredField("providers"); + field.setAccessible(true); + return (AbstractSubscriptionProviders) field.get(consumer); + } + private static class TestPushConsumer extends AbstractSubscriptionPushConsumer { private final List closedStatesDuringHandshake = new ArrayList<>(); @@ -136,7 +242,14 @@ protected AbstractSubscriptionProvider constructSubscriptionProvider( connectionTimeoutInMs, this::isClosed, closedStatesDuringHandshake, - closedStatesDuringClose); + closedStatesDuringClose, + () -> false, + () -> false, + () -> false, + () -> false, + () -> {}, + () -> {}, + () -> {}); } } @@ -144,6 +257,14 @@ private static class TestPullConsumer extends AbstractSubscriptionPullConsumer { private final List closedStatesDuringHandshake = new ArrayList<>(); private final List closedStatesDuringClose = new ArrayList<>(); + private final List 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 TestPullConsumer() { super( @@ -170,21 +291,31 @@ protected AbstractSubscriptionProvider constructSubscriptionProvider( 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); + 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++); + createdProviders.add(provider); + return provider; } } @@ -193,6 +324,13 @@ private static class TestSubscriptionProvider extends AbstractSubscriptionProvid private final BooleanSupplier consumerClosedSupplier; private final List closedStatesDuringHandshake; private final List 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 TestSubscriptionProvider( final TEndPoint endPoint, @@ -208,7 +346,14 @@ private TestSubscriptionProvider( final int connectionTimeoutInMs, final BooleanSupplier consumerClosedSupplier, final List closedStatesDuringHandshake, - final List closedStatesDuringClose) { + final List closedStatesDuringClose, + final BooleanSupplier fenceOnHandshake, + final BooleanSupplier fenceOnHeartbeat, + final BooleanSupplier fenceOnPollTablets, + final BooleanSupplier returnPartialTablets, + final Runnable closeRequest, + final Runnable commitRequest, + final Runnable sessionClose) { super( endPoint, username, @@ -224,6 +369,13 @@ private TestSubscriptionProvider( 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; } @Override @@ -249,19 +401,64 @@ protected AbstractSessionBuilder constructSubscriptionSessionBuilder( @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(); + setUnavailable(); + } + + @Override + synchronized void closeSession() { + closedStatesDuringClose.add(consumerClosedSupplier.getAsBoolean()); + sessionClose.run(); setUnavailable(); } @Override PipeSubscribeHeartbeatResp heartbeat( final List processorBufferedCommitContexts) { + if (fenceOnHeartbeat.getAsBoolean()) { + throw new SubscriptionConsumerFencedException("consumer connection fenced"); + } return new PipeSubscribeHeartbeatResp(); } + + @Override + List poll( + final Set topicNames, + final long timeoutMs, + final Map 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 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 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 0000000000000..5da15cadcb439 --- /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 b818a52e83877..d0236c1293dd3 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 @@ private 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 fe768a1f0181c..d11e16f41f64a 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 @@ private 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 3fdadb6e84e85..ad5540344c0ae 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.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> receiverConstructors = new HashMap<>(); private final ThreadLocal receiverThreadLocal = new ThreadLocal<>(); @@ -141,7 +151,7 @@ public TPipeSubscribeResp handle(final TPipeSubscribeReq req, final String usern if (isHandshake(req)) { if (isSuccessful(requestResult.response)) { if (currentReceiver != null && currentReceiver != receiver) { - currentReceiver.invalidateConsumer(); + invalidateReplacedReceiver(currentReceiver, identity); } return receiver; } @@ -158,7 +168,9 @@ public TPipeSubscribeResp handle(final TPipeSubscribeReq req, final String usern 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 +293,30 @@ private TPipeSubscribeResp handleRequest( 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; } - 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,6 +335,16 @@ private void removeReceiverMappings(final SubscriptionReceiver receiver) { (identity, currentReceiver) -> consumerReceivers.remove(identity, receiver)); } + 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 ConsumerIdentity getConsumerIdentity( final TPipeSubscribeReq req, final SubscriptionReceiver receiver) { if (isHandshake(req) && req.isSetBody()) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java index b0eec505a5c48..caa4b2b49fe6f 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 @@ -121,11 +121,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 consumerConfigThreadLocal = new ThreadLocal<>(); private final ThreadLocal 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 +145,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: @@ -206,7 +219,10 @@ public boolean hasActiveConsumer() { @Override public void invalidateConsumer() { - clearSharedConsumerState(); + synchronized (this) { + consumerFenced = true; + clearSharedConsumerState(); + } } @Override @@ -1340,17 +1356,23 @@ private void unsubscribe(final ConsumerConfig consumerConfig, final Set } } - 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 bdf75a3a33952..652eeca72aecc 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.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.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,257 @@ public void testReconnectInvalidatesOldReceiverBeforeTimeoutCleanup() throws IOE Assert.assertEquals(1, newReceiver.timeoutCount.get()); } + @Test + public void testDuplicateConnectionFencesOldReceiverWithoutInvalidatingNewReceiver() + throws Exception { + final CopyOnWriteArrayList 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 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 testConcurrentHandshakeWithSameIdentityFencesOldReceiver() throws Exception { + final CopyOnWriteArrayList 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 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 oldHandshakeResponse = new AtomicReference<>(); + final AtomicReference newHandshakeResponse = new AtomicReference<>(); + final AtomicReference oldHeartbeatResponse = new AtomicReference<>(); + final AtomicReference newHeartbeatResponse = new AtomicReference<>(); + final AtomicReference 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 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 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 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 oldHandshakeResponse = new AtomicReference<>(); + final AtomicReference newHandshakeResponse = new AtomicReference<>(); + final AtomicReference newHeartbeatResponse = new AtomicReference<>(); + final AtomicReference 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 receivers = new CopyOnWriteArrayList<>(); @@ -234,22 +488,53 @@ private TPipeSubscribeReq createHandshakeRequest( return PipeSubscribeHandshakeReq.toTPipeSubscribeReq(new ConsumerConfig(attributes)); } + 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 +544,7 @@ public TPipeSubscribeResp handle(final TPipeSubscribeReq req) { return response(TSStatusCode.SUCCESS_STATUS); } return response( - invalidated ? TSStatusCode.SUBSCRIPTION_MISSING_CONSUMER : TSStatusCode.SUCCESS_STATUS); + invalidated ? TSStatusCode.SUBSCRIPTION_CONSUMER_FENCED : TSStatusCode.SUCCESS_STATUS); } @Override 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 21d2e7b67d891..807d72315147b 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.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 void testHandleExitClearsThreadLocalStateAfterInvalidation() throws Excep 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(); From 10f6a89791b3fab381f098726f3d4e603fd8ee7a Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Thu, 17 Sep 2026 10:36:41 +0800 Subject: [PATCH 2/2] Make consumer fencing consistent across DataNodes --- .../subscription/config/ConsumerConfig.java | 4 + .../subscription/config/ConsumerConstant.java | 1 + .../base/AbstractSubscriptionConsumer.java | 117 +++++++++++------- .../base/AbstractSubscriptionProvider.java | 8 ++ .../SubscriptionConsumerLifecycleTest.java | 85 +++++++++++++ .../agent/SubscriptionReceiverAgent.java | 39 +++++- .../receiver/SubscriptionReceiver.java | 8 ++ .../receiver/SubscriptionReceiverV1.java | 12 +- .../agent/SubscriptionReceiverAgentTest.java | 70 +++++++++++ 9 files changed, 297 insertions(+), 47 deletions(-) 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 13f2a9ee3fb03..47926b5fc9112 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 String getConsumerGroupId() { 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 3df95facf367c..89878efb49a06 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/session/subscription/consumer/base/AbstractSubscriptionConsumer.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java index 15aaf6e43d324..9eb0ae809f22e 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 @@ -91,6 +91,7 @@ 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; @@ -109,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; @@ -118,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; @@ -197,6 +200,10 @@ public Long getOwnerEpoch() { return ownerEpoch; } + String getConsumerInstanceId() { + return consumerInstanceId; + } + /////////////////////////////// ctor /////////////////////////////// protected AbstractSubscriptionConsumer(final AbstractSubscriptionConsumerBuilder builder) { @@ -582,6 +589,7 @@ AbstractSubscriptionProvider constructProviderAndHandshake(final TEndPoint endPo this.thriftMaxFrameSize, this.heartbeatIntervalMs, this.connectionTimeoutInMs); + provider.setConsumerInstanceId(consumerInstanceId); try { provider.handshake(); } catch (final Exception e) { @@ -627,6 +635,13 @@ String sanitizeConnectionFailureMessage(final Throwable throwable) { 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 { @@ -784,42 +799,10 @@ protected List multiplePoll( tasks.add(new PollTask(partition, timeoutMs)); } - // submit multiple tasks to poll messages - final List messages = new ArrayList<>(); - SubscriptionRuntimeCriticalException lastSubscriptionRuntimeCriticalException = null; try { // strict timeout - for (final Future> 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 @@ -832,6 +815,56 @@ protected List multiplePoll( // TODO: ignore possible interrupted state? + return Collections.emptyList(); + } + + List collectMultiplePollResults( + final List>> futures, final Set topicNames) + throws InterruptedException { + final List messages = new ArrayList<>(); + SubscriptionRuntimeCriticalException lastSubscriptionRuntimeCriticalException = null; + for (final Future> 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)) { @@ -1711,11 +1744,11 @@ public AsyncCommitWorker( @Override public void run() { - if (isClosed()) { - return; - } - try { + checkIfFenced(); + if (isClosed()) { + return; + } ack(messages); callback.onComplete(); } catch (final Exception e) { @@ -1728,11 +1761,11 @@ protected CompletableFuture commitAsync(final Iterable future = new CompletableFuture<>(); SubscriptionExecutorServiceManager.submitAsyncCommitWorker( () -> { - if (isClosed()) { - return; - } - try { + checkIfFenced(); + if (isClosed()) { + return; + } ack(messages); future.complete(null); } catch (final Throwable e) { 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 2316f79aa95d4..c392f87771a69 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 @@ -88,6 +88,7 @@ public abstract class AbstractSubscriptionProvider { private String consumerId; private String consumerGroupId; + private String consumerInstanceId; private final String ownerId; private final Long ownerEpoch; @@ -174,6 +175,10 @@ String getConsumerGroupId() { return consumerGroupId; } + void setConsumerInstanceId(final String consumerInstanceId) { + this.consumerInstanceId = consumerInstanceId; + } + TEndPoint getEndPoint() { return endPoint; } @@ -191,6 +196,9 @@ synchronized void handshake() throws SubscriptionException, IoTDBConnectionExcep final Map 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); } 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 abc01f9875172..29a73e5a0925c 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 @@ -30,17 +30,26 @@ 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.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BooleanSupplier; public class SubscriptionConsumerLifecycleTest { @@ -191,6 +200,82 @@ public void testFencedTabletContinuationDoesNotSendNack() throws Exception { } } + @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> 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 messages = + () -> { + messagesIterated.set(true); + return Collections.emptyIterator(); + }; + + final CountDownLatch callbackCompleted = new CountDownLatch(1); + final AtomicBoolean callbackSucceeded = new AtomicBoolean(false); + final AtomicReference 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 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"); 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 ad5540344c0ae..dc42d4fbcae29 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 @@ -137,7 +137,8 @@ public TPipeSubscribeResp handle(final TPipeSubscribeReq req, final String usern 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)) { @@ -146,6 +147,15 @@ public TPipeSubscribeResp handle(final TPipeSubscribeReq req, final String usern 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)) { @@ -307,6 +317,12 @@ private boolean registerReceiver( registered.set(true); return receiver; } + if (receiver.getConsumerInstanceId() != null + && !shouldKeepCurrentReceiver(currentReceiver, receiver.getConsumerInstanceId())) { + invalidateReplacedReceiver(currentReceiver, key); + registered.set(true); + 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. @@ -345,7 +361,7 @@ private void invalidateReplacedReceiver( receiver.invalidateConsumer(); } - private static ConsumerIdentity getConsumerIdentity( + private static ConsumerConnection getConsumerConnection( final TPipeSubscribeReq req, final SubscriptionReceiver receiver) { if (isHandshake(req) && req.isSetBody()) { try { @@ -356,7 +372,7 @@ private static ConsumerIdentity getConsumerIdentity( ConsumerIdentity.of( consumerConfig.getConsumerGroupId(), consumerConfig.getConsumerId()); if (Objects.nonNull(identity)) { - return identity; + return new ConsumerConnection(identity, consumerConfig.getConsumerInstanceId()); } } } catch (final RuntimeException ignored) { @@ -364,7 +380,20 @@ private static ConsumerIdentity getConsumerIdentity( // 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) { @@ -396,4 +425,6 @@ private static ConsumerIdentity of(final String consumerGroupId, final String co : 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 e5de617c89c68..1d2fbe4fc1861 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 caa4b2b49fe6f..4c8c1c08d3047 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.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; @@ -212,6 +213,12 @@ public String getConsumerGroupId() { 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); @@ -1208,11 +1215,14 @@ private void unsubscribeCompleteTopics(final ConsumerConfig consumerConfig) { private void createConsumer(final ConsumerConfig consumerConfig) throws SubscriptionException { try (final ConfigNodeClient configNodeClient = CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) { + final Map 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( 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 652eeca72aecc..2d48518e5f22a 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 @@ -170,6 +170,40 @@ public void testDuplicateConnectionFencesOldReceiverWithoutInvalidatingNewReceiv 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 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 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 receivers = new CopyOnWriteArrayList<>(); @@ -482,12 +516,43 @@ private ScheduledExecutorService getReceiverTimeoutChecker(final SubscriptionRec 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 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 response = new AtomicReference<>(); + final AtomicReference 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<>())); } @@ -581,6 +646,11 @@ public String getConsumerGroupId() { return consumerConfig == null ? null : consumerConfig.getConsumerGroupId(); } + @Override + public String getConsumerInstanceId() { + return consumerConfig == null ? null : consumerConfig.getConsumerInstanceId(); + } + @Override public void invalidateConsumer() { consumerConfig = null;