This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new 495f04f7c17 Fix subscription provider recovery after broker loss
(#19018)
495f04f7c17 is described below
commit 495f04f7c179089aa261d19143c8fdd5b0601e94
Author: Caideyipi <[email protected]>
AuthorDate: Mon Oct 26 19:30:49 2026 +0800
Fix subscription provider recovery after broker loss (#19018)
* Fix subscription provider recovery after broker loss
* Preserve subscription state during provider recovery
Close stale sessions without dropping subscriptions, keep provider failures
terminal, ignore retired heartbeat results, and stop recovery after consumer
close. Recognize file and tablet broker-missing poll diagnostics in both server
locales and cover the recovery and commit paths with regression tests.
---
.../base/AbstractSubscriptionConsumer.java | 8 +
.../base/AbstractSubscriptionProvider.java | 55 +++-
.../base/AbstractSubscriptionProviders.java | 111 ++++---
.../SubscriptionConsumerMultiProviderPollTest.java | 336 +++++++++++++++++++--
.../SubscriptionProviderConnectionTimeoutTest.java | 130 ++++++++
5 files changed, 571 insertions(+), 69 deletions(-)
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 ea04b52109f..840fe5ac6ea 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
@@ -1453,6 +1453,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
if (isClosed()) {
return Collections.emptyList();
}
+ providers.syncIfUnavailable(this);
throw new SubscriptionConnectionException(
String.format(
SubscriptionMessages
@@ -1464,6 +1465,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
try {
return provider.pollFile(commitContext, writingOffset, timeoutMs);
} catch (final SubscriptionConnectionException ignored) {
+ providers.syncIfUnavailable(this);
return Collections.emptyList();
}
} finally {
@@ -1482,6 +1484,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
if (isClosed()) {
return Collections.emptyList();
}
+ providers.syncIfUnavailable(this);
throw new SubscriptionConnectionException(
String.format(
SubscriptionMessages
@@ -1493,6 +1496,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
try {
return provider.pollTablets(commitContext, offset, timeoutMs);
} catch (final SubscriptionConnectionException ignored) {
+ providers.syncIfUnavailable(this);
return Collections.emptyList();
}
} finally {
@@ -1684,6 +1688,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
if (isClosed()) {
return AbstractSubscriptionProvider.CommitResult.empty();
}
+ providers.syncIfUnavailable(this);
throw new SubscriptionConnectionException(
String.format(
SubscriptionMessages
@@ -1697,6 +1702,9 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
} catch (final SubscriptionConsumerFencedException e) {
fence(e);
throw e;
+ } catch (final SubscriptionConnectionException e) {
+ providers.syncIfUnavailable(this);
+ throw e;
}
} finally {
providers.releaseReadLock();
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 0beabf8d614..56f73249103 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
@@ -78,6 +78,7 @@ import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.ReentrantLock;
public abstract class AbstractSubscriptionProvider {
@@ -107,7 +108,15 @@ public abstract class AbstractSubscriptionProvider {
private final AtomicBoolean isClosed = new AtomicBoolean(true);
private final AtomicBoolean isClosing = new AtomicBoolean(false);
- private final AtomicBoolean isAvailable = new AtomicBoolean(false);
+
+ private enum Availability {
+ NEW,
+ AVAILABLE,
+ UNAVAILABLE
+ }
+
+ private final AtomicReference<Availability> availability =
+ new AtomicReference<>(Availability.NEW);
private final ReentrantLock rpcLock = new ReentrantLock();
private final int localThriftMaxFrameSize;
@@ -243,15 +252,17 @@ public abstract class AbstractSubscriptionProvider {
}
boolean isAvailable() {
- return isAvailable.get();
+ return availability.get() == Availability.AVAILABLE;
}
void setAvailable() {
- isAvailable.set(true);
+ // A heartbeat that completes after a failed RPC must not revive this
connection. Recovery
+ // constructs and handshakes a new provider instead.
+ availability.compareAndSet(Availability.NEW, Availability.AVAILABLE);
}
void setUnavailable() {
- isAvailable.set(false);
+ availability.set(Availability.UNAVAILABLE);
}
int getDataNodeId() {
@@ -413,10 +424,17 @@ public abstract class AbstractSubscriptionProvider {
}
synchronized void closeSession() throws IoTDBConnectionException {
+ setUnavailable();
try {
+ try {
+ // Recovery must preserve subscriptions and interrupt any heartbeat
still using this
+ // session.
+ getSessionConnection().forceClose();
+ } catch (final IoTDBConnectionException ignored) {
+ // A failed handshake may not have opened a session.
+ }
session.close();
} finally {
- setUnavailable();
isClosed.set(true);
}
}
@@ -736,11 +754,34 @@ public abstract class AbstractSubscriptionProvider {
setUnavailable();
throw new SubscriptionConnectionException(e.getMessage(), e);
}
- verifyPipeSubscribeSuccess(resp.status);
+ try {
+ verifyPipeSubscribeSuccess(resp.status);
+ } catch (final SubscriptionRuntimeNonCriticalException e) {
+ // A missing broker is reported as a poll error instead of a transport
failure. Treat it as a
+ // provider failure so the consumer can recreate this provider before
the next poll.
+ if (isMissingBrokerPollError(resp.status)) {
+ setUnavailable();
+ throw new SubscriptionConnectionException(e.getMessage(), e);
+ }
+ throw e;
+ }
final PipeSubscribePollResp pollResp =
PipeSubscribePollResp.fromTPipeSubscribeResp(resp);
return pollResp.getResponses();
}
+ private boolean isMissingBrokerPollError(final TSStatus status) {
+ if (status.code != TSStatusCode.SUBSCRIPTION_POLL_ERROR.getStatusCode()
+ || status.message == null) {
+ return false;
+ }
+ // These are server protocol diagnostics. Match both server locales
independently of the client
+ // locale, including the distinct messages returned by file and tablet
continuation polls.
+ return status.message.contains("no broker bound to consumer group [" +
consumerGroupId + "]")
+ || status.message.contains(
+ "broker bound to consumer group [" + consumerGroupId + "] does not
exist")
+ || status.message.contains("绑定到消费者组 [" + consumerGroupId + "] 的 Broker
不存在");
+ }
+
CommitResult commit(
final List<SubscriptionCommitContext> subscriptionCommitContexts, final
boolean nack)
throws SubscriptionException {
@@ -948,7 +989,7 @@ public abstract class AbstractSubscriptionProvider {
+ ", consumerGroupId="
+ consumerGroupId
+ ", isAvailable="
- + isAvailable
+ + isAvailable()
+ ", isClosed="
+ isClosed
+ "}";
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 565c719149d..afb5285b5f5 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
@@ -90,11 +90,14 @@ final class AbstractSubscriptionProviders {
/** Caller should ensure that the method is called in the lock {@link
#acquireWriteLock()}. */
void openProviders(final AbstractSubscriptionConsumer consumer) throws
SubscriptionException {
// close stale providers
- closeProviders();
+ closeProviders(false);
final Map<TEndPoint, String> connectionFailures = new LinkedHashMap<>();
final List<Throwable> connectionFailureCauses = new ArrayList<>();
for (final TEndPoint endPoint : initialEndpoints) {
+ if (isClosing.get() || consumer.isFenced()) {
+ return;
+ }
final AbstractSubscriptionProvider defaultProvider;
final int defaultDataNodeId;
@@ -116,6 +119,9 @@ final class AbstractSubscriptionProviders {
}
defaultDataNodeId = defaultProvider.getDataNodeId();
addProvider(defaultDataNodeId, defaultProvider);
+ if (isClosing.get() || consumer.isFenced()) {
+ return;
+ }
final Map<Integer, TEndPoint> allEndPoints;
try {
@@ -134,6 +140,9 @@ final class AbstractSubscriptionProviders {
}
for (final Map.Entry<Integer, TEndPoint> entry :
allEndPoints.entrySet()) {
+ if (isClosing.get() || consumer.isFenced()) {
+ return;
+ }
if (defaultDataNodeId == entry.getKey()) {
continue;
}
@@ -243,7 +252,8 @@ final class AbstractSubscriptionProviders {
}
final AbstractSubscriptionProvider provider =
subscriptionProviders.get(dataNodeId);
try {
- provider.close();
+ // Replacing a connection must not unsubscribe topics or drop the shared
consumer.
+ provider.closeSession();
} finally {
LOGGER.info(SubscriptionMessages.CLOSE_STALE_PROVIDER, provider);
subscriptionProviders.remove(dataNodeId);
@@ -338,13 +348,19 @@ final class AbstractSubscriptionProviders {
private void heartbeatProvider(
final AbstractSubscriptionConsumer consumer, final
AbstractSubscriptionProvider provider) {
try {
- if (consumer.isClosed() || consumer.isFenced() ||
!containsProvider(provider)) {
+ if (consumer.isClosed()
+ || consumer.isFenced()
+ || !containsProvider(provider)
+ || !provider.isAvailable()) {
return;
}
final List<SubscriptionCommitContext> processorBufferedCommitContexts =
consumer.getProcessorBufferedCommitContexts(provider.getDataNodeId());
final PipeSubscribeHeartbeatResp resp =
provider.heartbeat(processorBufferedCommitContexts);
- if (consumer.isClosed() || consumer.isFenced() ||
!containsProvider(provider)) {
+ if (consumer.isClosed()
+ || consumer.isFenced()
+ || !containsProvider(provider)
+ || !provider.isAvailable()) {
return;
}
provider.setAvailable();
@@ -360,7 +376,9 @@ final class AbstractSubscriptionProviders {
}
}
} catch (final SubscriptionConsumerFencedException e) {
- consumer.fence(e);
+ if (containsProvider(provider)) {
+ consumer.fence(e);
+ }
provider.setUnavailable();
} catch (final Exception e) {
LOGGER.warn(
@@ -398,7 +416,7 @@ final class AbstractSubscriptionProviders {
}
}
- /** Reconnects providers immediately after a failed RPC when every provider
is unavailable. */
+ /** Reconnects unavailable providers immediately after a failed RPC. */
void syncIfUnavailable(final AbstractSubscriptionConsumer consumer) {
if (consumer.isClosed()
|| consumer.isFenced()
@@ -415,7 +433,7 @@ final class AbstractSubscriptionProviders {
}
acquireWriteLock();
try {
- if (consumer.isClosed() || consumer.isFenced() ||
!hasNoAvailableProviders()) {
+ if (consumer.isClosed() || consumer.isFenced()) {
return;
}
final long now = System.currentTimeMillis();
@@ -436,7 +454,7 @@ final class AbstractSubscriptionProviders {
}
private void syncInternal(final AbstractSubscriptionConsumer consumer) {
- if (consumer.isFenced()) {
+ if (consumer.isClosed() || consumer.isFenced()) {
return;
}
if (hasNoAvailableProviders()) {
@@ -462,64 +480,71 @@ final class AbstractSubscriptionProviders {
return;
}
- // add new providers or handshake existing providers
+ // Add new providers, replace unavailable providers, or heartbeat existing
providers.
for (final Map.Entry<Integer, TEndPoint> entry : allEndPoints.entrySet()) {
- if (consumer.isFenced()) {
+ if (consumer.isClosed() || consumer.isFenced()) {
return;
}
- final AbstractSubscriptionProvider provider =
getProvider(entry.getKey());
- if (Objects.isNull(provider)) {
- // new provider
- final TEndPoint endPoint = entry.getValue();
- final AbstractSubscriptionProvider newProvider;
+ AbstractSubscriptionProvider provider = getProvider(entry.getKey());
+ if (Objects.nonNull(provider) && provider.isAvailable()) {
+ // Existing provider
try {
- newProvider = consumer.constructProviderAndHandshake(endPoint);
+ consumer.subscribedTopics = provider.heartbeat().getTopics();
+ provider.setAvailable();
} catch (final SubscriptionConsumerFencedException e) {
consumer.fence(e);
+ provider.setUnavailable();
return;
} catch (final Exception e) {
LOGGER.warn(
-
SubscriptionMessages.LOG_ARG_FAILED_CREATE_CONNECTION_ARG_BECAUSE_ARG_E536E22A,
+ SubscriptionMessages
+
.LOG_ARG_FAILED_SENDING_HEARTBEAT_SUBSCRIPTION_PROVIDER_ARG_BECAUSE_ARG_SET_0B38FB1F,
consumer,
- endPoint,
+ provider,
e,
e);
- continue;
+ provider.setUnavailable();
}
- addProvider(entry.getKey(), newProvider);
- } else {
- // existing provider
+ }
+
+ if (Objects.nonNull(provider) && !provider.isAvailable()) {
+ // Remove unavailable providers before creating replacements in this
same sync round.
try {
- consumer.subscribedTopics = provider.heartbeat().getTopics();
- provider.setAvailable();
- } catch (final SubscriptionConsumerFencedException e) {
- consumer.fence(e);
- provider.setUnavailable();
- return;
+ closeAndRemoveProvider(entry.getKey());
} catch (final Exception e) {
LOGGER.warn(
SubscriptionMessages
-
.LOG_ARG_FAILED_SENDING_HEARTBEAT_SUBSCRIPTION_PROVIDER_ARG_BECAUSE_ARG_SET_0B38FB1F,
+
.LOG_EXCEPTION_OCCURRED_ARG_CLOSING_REMOVING_SUBSCRIPTION_PROVIDER_ARG_BECAUSE_ARG_2EC38739,
consumer,
provider,
e,
e);
- provider.setUnavailable();
}
- // close and remove unavailable provider (reset the connection as much
as possible)
- if (!provider.isAvailable()) {
- try {
- closeAndRemoveProvider(entry.getKey());
- } catch (final Exception e) {
- LOGGER.warn(
- SubscriptionMessages
-
.LOG_EXCEPTION_OCCURRED_ARG_CLOSING_REMOVING_SUBSCRIPTION_PROVIDER_ARG_BECAUSE_ARG_2EC38739,
- consumer,
- provider,
- e,
- e);
- }
+ provider = null;
+ }
+
+ if (Objects.isNull(provider)) {
+ // New or replacement provider
+ if (consumer.isClosed() || consumer.isFenced()) {
+ return;
}
+ final TEndPoint endPoint = entry.getValue();
+ 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,
+ consumer,
+ endPoint,
+ e,
+ e);
+ continue;
+ }
+ addProvider(entry.getKey(), newProvider);
}
}
diff --git
a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerMultiProviderPollTest.java
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerMultiProviderPollTest.java
index 1d2ee7ae2c0..3ad5935b337 100644
---
a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerMultiProviderPollTest.java
+++
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionConsumerMultiProviderPollTest.java
@@ -22,6 +22,7 @@ package org.apache.iotdb.session.subscription.consumer.base;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.rpc.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.payload.poll.SubscriptionCommitContext;
import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollResponse;
@@ -40,6 +41,9 @@ import org.apache.tsfile.write.schema.MeasurementSchema;
import org.junit.Assert;
import org.junit.Test;
+import java.lang.reflect.Field;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
@@ -47,6 +51,12 @@ import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
@@ -134,11 +144,252 @@ public class SubscriptionConsumerMultiProviderPollTest {
}
}
+ @Test
+ public void testPollReconnectsWhenOneOfMultipleProvidersBecomesUnavailable()
+ throws SubscriptionException {
+ final TestPullConsumer consumer = new TestPullConsumer(1, 0, false, true,
3);
+ try {
+ consumer.open();
+ consumer.subscribeTopic();
+
+ final List<SubscriptionMessage> messages = consumer.pollForTest(2_000L);
+
+ Assert.assertEquals(1, messages.size());
+ Assert.assertEquals(4, consumer.getCreatedProviderCount());
+ Assert.assertEquals(2, consumer.getPollCount(1));
+ } finally {
+ consumer.close();
+ }
+ }
+
+ @Test
+ public void
testPartialRecoveryPreservesHealthyProvidersAndClosesOnlyFailedSession()
+ throws Exception {
+ final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+ try {
+ consumer.open();
+ final AbstractSubscriptionProviders providers = getProviders(consumer);
+ final TestSubscriptionProvider failed = consumer.createdProviders.get(0);
+ final AbstractSubscriptionProvider healthy = providers.getProvider(2);
+ failed.setUnavailable();
+
+ providers.sync(consumer);
+
+ Assert.assertEquals(4, consumer.getCreatedProviderCount());
+ Assert.assertNotSame(failed, providers.getProvider(1));
+ Assert.assertSame(healthy, providers.getProvider(2));
+ Assert.assertEquals(1, failed.sessionCloseCount.get());
+ Assert.assertEquals(0, failed.closeCount.get());
+ Assert.assertTrue(providers.getProvider(1).isAvailable());
+ } finally {
+ consumer.close();
+ }
+ }
+
+ @Test
+ public void testAllUnavailableRecoveryClosesSessionsWithoutClosingConsumer()
throws Exception {
+ final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+ try {
+ consumer.open();
+ final List<TestSubscriptionProvider> failedProviders =
+ new ArrayList<>(consumer.createdProviders);
+ failedProviders.forEach(AbstractSubscriptionProvider::setUnavailable);
+
+ getProviders(consumer).sync(consumer);
+
+ Assert.assertEquals(6, consumer.getCreatedProviderCount());
+ for (final TestSubscriptionProvider failed : failedProviders) {
+ Assert.assertEquals(1, failed.sessionCloseCount.get());
+ Assert.assertEquals(0, failed.closeCount.get());
+ }
+ } finally {
+ consumer.close();
+ }
+ }
+
+ @Test
+ public void testLateHeartbeatCannotReviveFailedProvider() throws Exception {
+ assertLateHeartbeatIgnored(false);
+ }
+
+ @Test
+ public void testRetiredHeartbeatCannotFenceRecoveredConsumer() throws
Exception {
+ assertLateHeartbeatIgnored(true);
+ }
+
+ private void assertLateHeartbeatIgnored(final boolean fenced) throws
Exception {
+ final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+ final ExecutorService executor = Executors.newSingleThreadExecutor();
+ final CountDownLatch heartbeatStarted = new CountDownLatch(1);
+ final CountDownLatch releaseHeartbeat = new CountDownLatch(1);
+ try {
+ consumer.open();
+ consumer.subscribeTopic();
+ final AbstractSubscriptionProviders providers = getProviders(consumer);
+ final TestSubscriptionProvider failed = consumer.createdProviders.get(0);
+ failed.emptyHeartbeatTopics = true;
+ failed.heartbeatAction =
+ () -> {
+ heartbeatStarted.countDown();
+ await(releaseHeartbeat);
+ if (fenced) {
+ throw new SubscriptionConsumerFencedException("retired
connection fenced");
+ }
+ };
+ final Method heartbeat =
+ AbstractSubscriptionProviders.class.getDeclaredMethod(
+ "heartbeatProvider",
+ AbstractSubscriptionConsumer.class,
+ AbstractSubscriptionProvider.class);
+ heartbeat.setAccessible(true);
+ final Future<?> result = executor.submit(() ->
heartbeat.invoke(providers, consumer, failed));
+ Assert.assertTrue(heartbeatStarted.await(5, TimeUnit.SECONDS));
+
+ failed.setUnavailable();
+ if (fenced) {
+ providers.sync(consumer);
+ }
+ releaseHeartbeat.countDown();
+ result.get(5, TimeUnit.SECONDS);
+
+ Assert.assertFalse(failed.isAvailable());
+ Assert.assertFalse(consumer.isFenced());
+ Assert.assertTrue(consumer.subscribedTopics.containsKey(TOPIC));
+ providers.sync(consumer);
+ Assert.assertNotSame(failed, providers.getProvider(1));
+ } finally {
+ releaseHeartbeat.countDown();
+ consumer.close();
+ executor.shutdownNow();
+ executor.awaitTermination(5, TimeUnit.SECONDS);
+ }
+ }
+
+ @Test
+ public void testCloseDuringReplacementHandshakeDoesNotOpenMoreProviders()
throws Exception {
+ assertCloseDuringReplacementHandshake(false);
+ assertCloseDuringReplacementHandshake(true);
+ }
+
+ private void assertCloseDuringReplacementHandshake(final boolean
allUnavailable)
+ throws Exception {
+ final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+ final ExecutorService executor = Executors.newSingleThreadExecutor();
+ final CountDownLatch handshakeStarted = new CountDownLatch(1);
+ final CountDownLatch releaseHandshake = new CountDownLatch(1);
+ try {
+ consumer.open();
+ final AbstractSubscriptionProviders providers = getProviders(consumer);
+ if (allUnavailable) {
+
consumer.createdProviders.forEach(AbstractSubscriptionProvider::setUnavailable);
+ } else {
+ consumer.createdProviders.get(0).setUnavailable();
+ }
+ consumer.handshakeAction =
+ () -> {
+ handshakeStarted.countDown();
+ await(releaseHandshake);
+ };
+ final Future<?> sync = executor.submit(() -> providers.sync(consumer));
+ Assert.assertTrue(handshakeStarted.await(5, TimeUnit.SECONDS));
+
+ consumer.close();
+ releaseHandshake.countDown();
+ sync.get(5, TimeUnit.SECONDS);
+
+ Assert.assertTrue(consumer.isClosed());
+ Assert.assertTrue(providers.hasNoProviders());
+ Assert.assertEquals(4, consumer.getCreatedProviderCount());
+ Assert.assertEquals(1,
consumer.createdProviders.get(3).closeCount.get());
+ } finally {
+ releaseHandshake.countDown();
+ consumer.close();
+ executor.shutdownNow();
+ executor.awaitTermination(5, TimeUnit.SECONDS);
+ }
+ }
+
+ @Test
+ public void testContinuationAndCommitFailuresSchedulePartialRecovery()
throws Exception {
+ for (final String operation :
+ new String[] {"pollFileInternal", "pollTabletsInternal",
"commitInternal"}) {
+ final TestPullConsumer consumer = new TestPullConsumer(1, 0);
+ try {
+ consumer.open();
+ final TestSubscriptionProvider failed =
consumer.createdProviders.get(0);
+ failed.failContinuation = true;
+ final CountDownLatch replacementStarted = new CountDownLatch(1);
+ consumer.handshakeAction = replacementStarted::countDown;
+ final SubscriptionCommitContext context =
+ new SubscriptionCommitContext(1, 0, TOPIC, CONSUMER_GROUP_ID, 0L);
+ final Method method;
+ final Object[] arguments;
+ if ("commitInternal".equals(operation)) {
+ method =
+ AbstractSubscriptionConsumer.class.getDeclaredMethod(
+ operation, int.class, List.class, boolean.class);
+ arguments = new Object[] {1, Collections.singletonList(context),
false};
+ } else {
+ method =
+ AbstractSubscriptionConsumer.class.getDeclaredMethod(
+ operation,
+ SubscriptionCommitContext.class,
+ "pollFileInternal".equals(operation) ? long.class :
int.class,
+ long.class);
+ arguments =
+ new Object[] {
+ context, "pollFileInternal".equals(operation) ? (Object) 0L :
0, 1_000L
+ };
+ }
+ method.setAccessible(true);
+ try {
+ method.invoke(consumer, arguments);
+ Assert.assertNotEquals("commitInternal", operation);
+ } catch (final InvocationTargetException e) {
+ Assert.assertEquals("commitInternal", operation);
+ Assert.assertTrue(e.getCause() instanceof
SubscriptionConnectionException);
+ }
+
+ Assert.assertTrue(replacementStarted.await(5, TimeUnit.SECONDS));
+ final AbstractSubscriptionProviders providers = getProviders(consumer);
+ providers.acquireReadLock();
+ try {
+ Assert.assertNotSame(failed, providers.getProvider(1));
+ Assert.assertEquals(3, providers.getAvailableProviderCount());
+ Assert.assertEquals("commitInternal".equals(operation) ? 1 : 0,
failed.commitCount.get());
+ Assert.assertEquals(0,
consumer.createdProviders.get(3).commitCount.get());
+ } finally {
+ providers.releaseReadLock();
+ }
+ } finally {
+ consumer.close();
+ }
+ }
+ }
+
+ private static AbstractSubscriptionProviders getProviders(
+ final AbstractSubscriptionConsumer consumer) throws Exception {
+ final Field field =
AbstractSubscriptionConsumer.class.getDeclaredField("providers");
+ field.setAccessible(true);
+ return (AbstractSubscriptionProviders) field.get(consumer);
+ }
+
+ private static void await(final CountDownLatch latch) {
+ try {
+ Assert.assertTrue(latch.await(5, TimeUnit.SECONDS));
+ } catch (final InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new AssertionError(e);
+ }
+ }
+
private static class TestPullConsumer extends
AbstractSubscriptionPullConsumer {
private final Map<Integer, Integer> pollCounts = new HashMap<>();
private final List<SubscriptionCommitContext> acknowledgedCommitContexts =
new ArrayList<>();
private final AtomicInteger createdProviderCount = new AtomicInteger();
+ private final List<TestSubscriptionProvider> createdProviders = new
CopyOnWriteArrayList<>();
+ private volatile Runnable handshakeAction = () -> {};
private final AtomicBoolean failFirstPoll;
private final int dataProviderId;
private final int emptyPollsBeforeData;
@@ -191,24 +442,28 @@ public class SubscriptionConsumerMultiProviderPollTest {
final long heartbeatIntervalMs,
final int connectionTimeoutInMs) {
createdProviderCount.incrementAndGet();
- return new TestSubscriptionProvider(
- endPoint,
- username,
- password,
- encryptedPassword,
- consumerId,
- consumerGroupId,
- ownerId,
- ownerEpoch,
- thriftMaxFrameSize,
- heartbeatIntervalMs,
- connectionTimeoutInMs,
- pollCounts,
- dataProviderId,
- emptyPollsBeforeData,
- emptyTabletBatch,
- failFirstPoll,
- providerCount);
+ final TestSubscriptionProvider provider =
+ new TestSubscriptionProvider(
+ endPoint,
+ username,
+ password,
+ encryptedPassword,
+ consumerId,
+ consumerGroupId,
+ ownerId,
+ ownerEpoch,
+ thriftMaxFrameSize,
+ heartbeatIntervalMs,
+ connectionTimeoutInMs,
+ pollCounts,
+ dataProviderId,
+ emptyPollsBeforeData,
+ emptyTabletBatch,
+ failFirstPoll,
+ providerCount);
+ provider.handshakeAction = () -> handshakeAction.run();
+ createdProviders.add(provider);
+ return provider;
}
private void subscribeTopic() {
@@ -251,6 +506,14 @@ public class SubscriptionConsumerMultiProviderPollTest {
private static class TestSubscriptionProvider extends
AbstractSubscriptionProvider {
+ private final AtomicInteger closeCount = new AtomicInteger();
+ private final AtomicInteger sessionCloseCount = new AtomicInteger();
+ private final AtomicInteger commitCount = new AtomicInteger();
+ private volatile Runnable heartbeatAction = () -> {};
+ private boolean emptyHeartbeatTopics;
+ private Runnable handshakeAction = () -> {};
+ private boolean failContinuation;
+
private final int dataNodeId;
private final Map<Integer, Integer> pollCounts;
private final int dataProviderId;
@@ -320,11 +583,19 @@ public class SubscriptionConsumerMultiProviderPollTest {
@Override
synchronized void handshake() {
+ handshakeAction.run();
setAvailable();
}
@Override
synchronized void close() {
+ closeCount.incrementAndGet();
+ setUnavailable();
+ }
+
+ @Override
+ synchronized void closeSession() {
+ sessionCloseCount.incrementAndGet();
setUnavailable();
}
@@ -336,8 +607,11 @@ public class SubscriptionConsumerMultiProviderPollTest {
@Override
PipeSubscribeHeartbeatResp heartbeat(
final List<SubscriptionCommitContext> processorBufferedCommitContexts)
{
+ heartbeatAction.run();
final PipeSubscribeHeartbeatResp response = new
PipeSubscribeHeartbeatResp();
- response.getTopics().put(TOPIC, new TopicConfig());
+ if (!emptyHeartbeatTopics) {
+ response.getTopics().put(TOPIC, new TopicConfig());
+ }
for (int id = 1; id <= providerCount; id++) {
response.getEndPoints().put(id, new TEndPoint(HOST, FIRST_PORT + id -
1));
}
@@ -389,6 +663,7 @@ public class SubscriptionConsumerMultiProviderPollTest {
List<SubscriptionPollResponse> pollTablets(
final SubscriptionCommitContext commitContext, final int offset, final
long timeoutMs)
throws SubscriptionException {
+ failContinuationIfRequested();
if (emptyTabletBatch && dataNodeId == dataProviderId && offset == 1) {
return Collections.singletonList(
new SubscriptionPollResponse(
@@ -398,5 +673,28 @@ public class SubscriptionConsumerMultiProviderPollTest {
}
return Collections.emptyList();
}
+
+ @Override
+ List<SubscriptionPollResponse> pollFile(
+ final SubscriptionCommitContext commitContext,
+ final long writingOffset,
+ final long timeoutMs) {
+ failContinuationIfRequested();
+ return Collections.emptyList();
+ }
+
+ @Override
+ CommitResult commit(final List<SubscriptionCommitContext> contexts, final
boolean nack) {
+ commitCount.incrementAndGet();
+ failContinuationIfRequested();
+ return new CommitResult(contexts, Collections.emptyMap());
+ }
+
+ private void failContinuationIfRequested() {
+ if (failContinuation) {
+ setUnavailable();
+ throw new SubscriptionConnectionException("transport failed");
+ }
+ }
}
}
diff --git
a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionProviderConnectionTimeoutTest.java
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionProviderConnectionTimeoutTest.java
index ab7a741b07d..46dbd1f8e9f 100644
---
a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionProviderConnectionTimeoutTest.java
+++
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionProviderConnectionTimeoutTest.java
@@ -22,10 +22,13 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
import org.apache.iotdb.rpc.TSStatusCode;
import
org.apache.iotdb.rpc.subscription.exception.SubscriptionConnectionException;
+import
org.apache.iotdb.rpc.subscription.exception.SubscriptionRuntimeNonCriticalException;
import
org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext;
+import
org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeRequestType;
import
org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeHandshakeResp;
import org.apache.iotdb.service.rpc.thrift.IClientRPCService;
import org.apache.iotdb.service.rpc.thrift.ServerProperties;
+import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeReq;
import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeResp;
import org.apache.iotdb.service.rpc.thrift.TSOpenSessionResp;
import org.apache.iotdb.service.rpc.thrift.TSProtocolVersion;
@@ -128,6 +131,133 @@ public class SubscriptionProviderConnectionTimeoutTest {
}
}
+ @Test
+ public void testMissingBrokerPollErrorMarksProviderUnavailable() throws
Exception {
+ assertPollError(
+ "Subscription: no broker bound to consumer group [group]", true /*
broker missing */);
+ }
+
+ @Test
+ public void testMissingFileBrokerPollErrorMarksProviderUnavailable() throws
Exception {
+ assertPollError(
+ "Subscription: pipe broker bound to consumer group [group] does not
exist", true);
+ }
+
+ @Test
+ public void testMissingTabletBrokerPollErrorMarksProviderUnavailable()
throws Exception {
+ assertPollError("Subscription: broker bound to consumer group [group] does
not exist", true);
+ }
+
+ @Test
+ public void
testMissingTabletBrokerPollErrorFromChineseServerMarksProviderUnavailable()
+ throws Exception {
+ assertPollError("订阅：绑定到消费者组 [group] 的 Broker 不存在", true);
+ }
+
+ @Test
+ public void testOtherPollErrorRemainsNonCritical() throws Exception {
+ assertPollError("other poll error", false /* broker missing */);
+ }
+
+ @Test
+ public void testNullPollErrorRemainsNonCritical() throws Exception {
+ assertPollError(null, false /* broker missing */);
+ }
+
+ private void assertPollError(final String message, final boolean
brokerMissing) throws Exception {
+ final AtomicInteger subscriptionCloseCount = new AtomicInteger();
+ final AtomicInteger sessionCloseCount = new AtomicInteger();
+ final ServerSocket socket = new ServerSocket(0, 50,
InetAddress.getByName("127.0.0.1"));
+ final TServerSocket serverSocket = new TServerSocket(socket);
+ final IClientRPCService.Iface service =
+ (IClientRPCService.Iface)
+ Proxy.newProxyInstance(
+ IClientRPCService.Iface.class.getClassLoader(),
+ new Class<?>[] {IClientRPCService.Iface.class},
+ (proxy, method, args) -> {
+ switch (method.getName()) {
+ case "openSession":
+ return new TSOpenSessionResp(
+ successStatus(),
TSProtocolVersion.IOTDB_SERVICE_PROTOCOL_V3)
+ .setSessionId(1L);
+ case "requestStatementId":
+ return 1L;
+ case "getProperties":
+ return new ServerProperties(
+ "test", Arrays.asList("max_time", "min_time"), "ms",
1);
+ case "pipeSubscribe":
+ final TPipeSubscribeReq request = (TPipeSubscribeReq)
args[0];
+ if (request.getType() ==
PipeSubscribeRequestType.HANDSHAKE.getType()) {
+ return PipeSubscribeHandshakeResp.toTPipeSubscribeResp(
+ successStatus(), 1, "consumer", "group");
+ }
+ if (request.getType() !=
PipeSubscribeRequestType.POLL.getType()) {
+ if (request.getType() ==
PipeSubscribeRequestType.CLOSE.getType()) {
+ subscriptionCloseCount.incrementAndGet();
+ }
+ return new TPipeSubscribeResp(successStatus(), (byte)
1, (short) 6);
+ }
+ return new TPipeSubscribeResp(
+ new
TSStatus(TSStatusCode.SUBSCRIPTION_POLL_ERROR.getStatusCode())
+ .setMessage(message),
+ (byte) 1,
+ (short) 6);
+ case "closeSession":
+ sessionCloseCount.incrementAndGet();
+ return successStatus();
+ default:
+ return defaultValue(method.getReturnType());
+ }
+ });
+ final TSimpleServer server =
+ new TSimpleServer(
+ new TSimpleServer.Args(serverSocket)
+ .processor(new IClientRPCService.Processor<>(service))
+ .transportFactory(new
org.apache.iotdb.rpc.TElasticFramedTransport.Factory()));
+ final Thread serverThread =
+ new Thread(server::serve, "subscription-provider-poll-error-server");
+ serverThread.setDaemon(true);
+ serverThread.start();
+
+ final TestSubscriptionProvider provider = new
TestSubscriptionProvider(socket.getLocalPort());
+ try {
+ provider.handshake();
+ final TTransport transport =
getTransport(provider.getSessionConnection());
+ final SubscriptionCommitContext context =
+ new SubscriptionCommitContext(1, 0, "topic", "group", 1L);
+
+ if (brokerMissing) {
+ Assert.assertThrows(
+ SubscriptionConnectionException.class,
+ () -> provider.poll(Collections.singleton("topic"), 1_000L));
+ Assert.assertFalse(provider.isAvailable());
+ Assert.assertThrows(
+ SubscriptionConnectionException.class, () ->
provider.pollFile(context, 0L, 1_000L));
+ Assert.assertThrows(
+ SubscriptionConnectionException.class, () ->
provider.pollTablets(context, 0, 1_000L));
+ provider.closeSession();
+ Assert.assertFalse(transport.isOpen());
+ Assert.assertEquals(0, subscriptionCloseCount.get());
+ Assert.assertEquals(0, sessionCloseCount.get());
+ } else {
+ Assert.assertThrows(
+ SubscriptionRuntimeNonCriticalException.class,
+ () -> provider.poll(Collections.singleton("topic"), 1_000L));
+ Assert.assertTrue(provider.isAvailable());
+ Assert.assertThrows(
+ SubscriptionRuntimeNonCriticalException.class,
+ () -> provider.pollFile(context, 0L, 1_000L));
+ Assert.assertThrows(
+ SubscriptionRuntimeNonCriticalException.class,
+ () -> provider.pollTablets(context, 0, 1_000L));
+ }
+ } finally {
+ provider.close();
+ server.stop();
+ serverThread.join(2_000);
+ }
+ }
+
private static TTransport getTransport(final SubscriptionSessionConnection
connection)
throws Exception {
final Field field = SessionConnection.class.getDeclaredField("transport");
