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 1e851a0e5ee Fix subscription consumer recovery after RPC timeout
(#18980)
1e851a0e5ee is described below
commit 1e851a0e5eefb09430619c1e14345b56f61b8d4f
Author: Caideyipi <[email protected]>
AuthorDate: Thu Oct 22 12:24:20 2026 +0800
Fix subscription consumer recovery after RPC timeout (#18980)
---
.../base/AbstractSubscriptionConsumer.java | 26 +--
.../base/AbstractSubscriptionProvider.java | 9 +-
.../base/AbstractSubscriptionProviders.java | 40 ++++
.../SubscriptionConsumerMultiProviderPollTest.java | 63 ++++++-
.../SubscriptionProviderConnectionTimeoutTest.java | 207 +++++++++++++++++++++
5 files changed, 326 insertions(+), 19 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 1cb614c6546..ea04b52109f 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
@@ -1415,6 +1415,7 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
private List<SubscriptionPollResponse> pollInternal(
final Set<String> topicNames, final long timeoutMs) throws
SubscriptionException {
checkIfFenced();
+ boolean syncIfUnavailable = false;
providers.acquireReadLock();
try {
final AbstractSubscriptionProvider provider =
providers.getNextAvailableProvider();
@@ -1422,22 +1423,23 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
if (isClosed()) {
return Collections.emptyList();
}
- throw new SubscriptionConnectionException(
- String.format(
- SubscriptionMessages
-
.EXCEPTION_CLUSTER_HAS_NO_AVAILABLE_SUBSCRIPTION_PROVIDERS_ARG_POLL_TOPIC_ARG_FABF7A4E,
- this,
- topicNames));
- }
- // ignore SubscriptionConnectionException to improve poll auto retry
- try {
- return provider.poll(topicNames, timeoutMs,
buildCurrentProgressByTopic(topicNames));
- } catch (final SubscriptionConnectionException ignored) {
- return Collections.emptyList();
+ syncIfUnavailable = true;
+ } else {
+ // ignore SubscriptionConnectionException to improve poll auto retry
+ try {
+ return provider.poll(topicNames, timeoutMs,
buildCurrentProgressByTopic(topicNames));
+ } catch (final SubscriptionConnectionException ignored) {
+ syncIfUnavailable = true;
+ }
}
} finally {
providers.releaseReadLock();
}
+
+ if (syncIfUnavailable) {
+ providers.syncIfUnavailable(this);
+ }
+ return Collections.emptyList();
}
private List<SubscriptionPollResponse> pollFileInternal(
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 db0ded818ec..0beabf8d614 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
@@ -174,9 +174,16 @@ public abstract class AbstractSubscriptionProvider {
private TPipeSubscribeResp pipeSubscribe(final TPipeSubscribeReq req)
throws TException, IoTDBConnectionException {
rpcLock.lock();
+ SubscriptionSessionConnection connection = null;
try {
- final SubscriptionSessionConnection connection = getSessionConnection();
+ connection = getSessionConnection();
return pipeSubscribeWithOptionalSlicing(req, requestBodySizeLimit,
connection::pipeSubscribe);
+ } catch (final TException | IoTDBConnectionException e) {
+ setUnavailable();
+ if (connection != null) {
+ connection.forceClose();
+ }
+ throw e;
} finally {
rpcLock.unlock();
}
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 cae7ded2cf9..565c719149d 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
@@ -21,6 +21,7 @@ package org.apache.iotdb.session.subscription.consumer.base;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.rpc.IoTDBConnectionException;
+import org.apache.iotdb.rpc.subscription.config.ConsumerConstant;
import
org.apache.iotdb.rpc.subscription.exception.SubscriptionConnectionException;
import
org.apache.iotdb.rpc.subscription.exception.SubscriptionConsumerFencedException;
import org.apache.iotdb.rpc.subscription.exception.SubscriptionException;
@@ -51,6 +52,7 @@ final class AbstractSubscriptionProviders {
private final SortedMap<Integer, AbstractSubscriptionProvider>
subscriptionProviders =
new ConcurrentSkipListMap<>();
private final AtomicBoolean isClosing = new AtomicBoolean(false);
+ private final AtomicBoolean unavailableSyncInFlight = new
AtomicBoolean(false);
private final Set<AbstractSubscriptionProvider>
providersWithHeartbeatInFlight =
ConcurrentHashMap.newKeySet();
private final Object heartbeatResponseLock = new Object();
@@ -59,6 +61,7 @@ final class AbstractSubscriptionProviders {
private final ReentrantReadWriteLock subscriptionProvidersLock = new
ReentrantReadWriteLock(true);
private final Set<TEndPoint> initialEndpoints;
+ private long lastUnavailableSyncTimeMs;
AbstractSubscriptionProviders(final Set<TEndPoint> initialEndpoints) {
this.initialEndpoints = initialEndpoints;
@@ -395,6 +398,43 @@ final class AbstractSubscriptionProviders {
}
}
+ /** Reconnects providers immediately after a failed RPC when every provider
is unavailable. */
+ void syncIfUnavailable(final AbstractSubscriptionConsumer consumer) {
+ if (consumer.isClosed()
+ || consumer.isFenced()
+ || !unavailableSyncInFlight.compareAndSet(false, true)) {
+ return;
+ }
+
+ if (Objects.isNull(
+ SubscriptionExecutorServiceManager.submitProviderHeartbeat(
+ () -> {
+ try {
+ if (consumer.isClosed() || consumer.isFenced()) {
+ return;
+ }
+ acquireWriteLock();
+ try {
+ if (consumer.isClosed() || consumer.isFenced() ||
!hasNoAvailableProviders()) {
+ return;
+ }
+ final long now = System.currentTimeMillis();
+ if (now - lastUnavailableSyncTimeMs
+ >=
ConsumerConstant.ENDPOINTS_SYNC_INTERVAL_MS_MIN_VALUE) {
+ lastUnavailableSyncTimeMs = now;
+ syncInternal(consumer);
+ }
+ } finally {
+ releaseWriteLock();
+ }
+ } finally {
+ unavailableSyncInFlight.set(false);
+ }
+ }))) {
+ unavailableSyncInFlight.set(false);
+ }
+ }
+
private void syncInternal(final AbstractSubscriptionConsumer consumer) {
if (consumer.isFenced()) {
return;
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 5faeff91ac4..1d2ee7ae2c0 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
@@ -21,6 +21,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.SubscriptionException;
import
org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext;
import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollResponse;
@@ -46,6 +47,8 @@ import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
public class SubscriptionConsumerMultiProviderPollTest {
@@ -114,21 +117,50 @@ public class SubscriptionConsumerMultiProviderPollTest {
}
}
+ @Test
+ public void testPollReconnectsWhenProviderBecomesUnavailable() throws
SubscriptionException {
+ final TestPullConsumer consumer = new TestPullConsumer(1, 0, false, true,
1);
+ try {
+ consumer.open();
+ consumer.subscribeTopic();
+
+ final List<SubscriptionMessage> messages = consumer.pollForTest(1_000L);
+
+ Assert.assertEquals(1, messages.size());
+ Assert.assertEquals(2, consumer.getCreatedProviderCount());
+ Assert.assertEquals(2, consumer.getPollCount(1));
+ } finally {
+ consumer.close();
+ }
+ }
+
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 AtomicBoolean failFirstPoll;
private final int dataProviderId;
private final int emptyPollsBeforeData;
private final boolean emptyTabletBatch;
+ private final int providerCount;
private int backoffCount;
private TestPullConsumer(final int dataProviderId, final int
emptyPollsBeforeData) {
- this(dataProviderId, emptyPollsBeforeData, false);
+ this(dataProviderId, emptyPollsBeforeData, false, false, 3);
}
private TestPullConsumer(
final int dataProviderId, final int emptyPollsBeforeData, final
boolean emptyTabletBatch) {
+ this(dataProviderId, emptyPollsBeforeData, emptyTabletBatch, false, 3);
+ }
+
+ private TestPullConsumer(
+ final int dataProviderId,
+ final int emptyPollsBeforeData,
+ final boolean emptyTabletBatch,
+ final boolean failFirstPoll,
+ final int providerCount) {
super(
new AbstractSubscriptionPullConsumerBuilder()
.host(HOST)
@@ -141,6 +173,8 @@ public class SubscriptionConsumerMultiProviderPollTest {
this.dataProviderId = dataProviderId;
this.emptyPollsBeforeData = emptyPollsBeforeData;
this.emptyTabletBatch = emptyTabletBatch;
+ this.failFirstPoll = new AtomicBoolean(failFirstPoll);
+ this.providerCount = providerCount;
}
@Override
@@ -156,6 +190,7 @@ public class SubscriptionConsumerMultiProviderPollTest {
final int thriftMaxFrameSize,
final long heartbeatIntervalMs,
final int connectionTimeoutInMs) {
+ createdProviderCount.incrementAndGet();
return new TestSubscriptionProvider(
endPoint,
username,
@@ -171,7 +206,9 @@ public class SubscriptionConsumerMultiProviderPollTest {
pollCounts,
dataProviderId,
emptyPollsBeforeData,
- emptyTabletBatch);
+ emptyTabletBatch,
+ failFirstPoll,
+ providerCount);
}
private void subscribeTopic() {
@@ -195,6 +232,10 @@ public class SubscriptionConsumerMultiProviderPollTest {
return acknowledgedCommitContexts;
}
+ private int getCreatedProviderCount() {
+ return createdProviderCount.get();
+ }
+
@Override
protected void ack(final Iterable<SubscriptionMessage> messages) throws
SubscriptionException {
for (final SubscriptionMessage message : messages) {
@@ -215,6 +256,8 @@ public class SubscriptionConsumerMultiProviderPollTest {
private final int dataProviderId;
private final int emptyPollsBeforeData;
private final boolean emptyTabletBatch;
+ private final AtomicBoolean failFirstPoll;
+ private final int providerCount;
private TestSubscriptionProvider(
final TEndPoint endPoint,
@@ -231,7 +274,9 @@ public class SubscriptionConsumerMultiProviderPollTest {
final Map<Integer, Integer> pollCounts,
final int dataProviderId,
final int emptyPollsBeforeData,
- final boolean emptyTabletBatch) {
+ final boolean emptyTabletBatch,
+ final AtomicBoolean failFirstPoll,
+ final int providerCount) {
super(
endPoint,
username,
@@ -249,6 +294,8 @@ public class SubscriptionConsumerMultiProviderPollTest {
this.dataProviderId = dataProviderId;
this.emptyPollsBeforeData = emptyPollsBeforeData;
this.emptyTabletBatch = emptyTabletBatch;
+ this.failFirstPoll = failFirstPoll;
+ this.providerCount = providerCount;
}
@Override
@@ -291,9 +338,9 @@ public class SubscriptionConsumerMultiProviderPollTest {
final List<SubscriptionCommitContext> processorBufferedCommitContexts)
{
final PipeSubscribeHeartbeatResp response = new
PipeSubscribeHeartbeatResp();
response.getTopics().put(TOPIC, new TopicConfig());
- response.getEndPoints().put(1, new TEndPoint(HOST, FIRST_PORT));
- response.getEndPoints().put(2, new TEndPoint(HOST, FIRST_PORT + 1));
- response.getEndPoints().put(3, new TEndPoint(HOST, FIRST_PORT + 2));
+ for (int id = 1; id <= providerCount; id++) {
+ response.getEndPoints().put(id, new TEndPoint(HOST, FIRST_PORT + id -
1));
+ }
return response;
}
@@ -304,6 +351,10 @@ public class SubscriptionConsumerMultiProviderPollTest {
final Map<String, TopicProgress> progressByTopic)
throws SubscriptionException {
pollCounts.merge(dataNodeId, 1, Integer::sum);
+ if (dataNodeId == dataProviderId && failFirstPoll.compareAndSet(true,
false)) {
+ setUnavailable();
+ throw new SubscriptionConnectionException("poll timed out");
+ }
if (dataNodeId != dataProviderId || pollCounts.get(dataNodeId) <=
emptyPollsBeforeData) {
return Collections.emptyList();
}
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
new file mode 100644
index 00000000000..ab7a741b07d
--- /dev/null
+++
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionProviderConnectionTimeoutTest.java
@@ -0,0 +1,207 @@
+/*
+ * 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.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.payload.poll.SubscriptionCommitContext;
+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.TPipeSubscribeResp;
+import org.apache.iotdb.service.rpc.thrift.TSOpenSessionResp;
+import org.apache.iotdb.service.rpc.thrift.TSProtocolVersion;
+import org.apache.iotdb.session.AbstractSessionBuilder;
+import org.apache.iotdb.session.SessionConnection;
+import org.apache.iotdb.session.subscription.SubscriptionSessionConnection;
+import org.apache.iotdb.session.subscription.SubscriptionTreeSessionBuilder;
+
+import org.apache.thrift.server.TSimpleServer;
+import org.apache.thrift.transport.TServerSocket;
+import org.apache.thrift.transport.TTransport;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.lang.reflect.Field;
+import java.lang.reflect.Proxy;
+import java.net.InetAddress;
+import java.net.ServerSocket;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Objects;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class SubscriptionProviderConnectionTimeoutTest {
+
+ private static final int READ_TIMEOUT_MS = 200;
+ private static final int CONNECTION_TIMEOUT_MS = 5_000;
+
+ @Test(timeout = 10_000)
+ public void testCommitReadTimeoutClosesProviderTransport() throws Exception {
+ final CountDownLatch commitStarted = new CountDownLatch(1);
+ final CountDownLatch releaseCommit = new CountDownLatch(1);
+ final AtomicInteger pipeSubscribeCount = 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":
+ if (pipeSubscribeCount.getAndIncrement() == 0) {
+ return PipeSubscribeHandshakeResp.toTPipeSubscribeResp(
+ successStatus(), 1, "consumer", "group");
+ }
+ commitStarted.countDown();
+ releaseCommit.await(5, TimeUnit.SECONDS);
+ return new TPipeSubscribeResp(successStatus(), (byte) 1,
(short) 6);
+ case "closeSession":
+ 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-timeout-server");
+ serverThread.setDaemon(true);
+ serverThread.start();
+
+ final TestSubscriptionProvider provider = new
TestSubscriptionProvider(socket.getLocalPort());
+ try {
+ provider.handshake();
+ final TTransport transport =
getTransport(provider.getSessionConnection());
+ Assert.assertTrue(transport.isOpen());
+
Assert.assertTrue(provider.getSessionConnection().setTimeout(READ_TIMEOUT_MS));
+
+ Assert.assertThrows(
+ SubscriptionConnectionException.class,
+ () ->
+ provider.commit(
+ Collections.singletonList(
+ new SubscriptionCommitContext(1, 0, "topic", "group",
1L)),
+ false));
+
+ Assert.assertTrue(commitStarted.await(1, TimeUnit.SECONDS));
+ Assert.assertFalse(provider.isAvailable());
+ Assert.assertFalse(transport.isOpen());
+ } finally {
+ provider.close();
+ releaseCommit.countDown();
+ server.stop();
+ serverThread.join(2_000);
+ }
+ }
+
+ private static TTransport getTransport(final SubscriptionSessionConnection
connection)
+ throws Exception {
+ final Field field = SessionConnection.class.getDeclaredField("transport");
+ field.setAccessible(true);
+ return (TTransport) field.get(connection);
+ }
+
+ private static TSStatus successStatus() {
+ return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
+ }
+
+ private static Object defaultValue(final Class<?> returnType) {
+ if (!returnType.isPrimitive() || returnType == Void.TYPE) {
+ return null;
+ }
+ if (returnType == Boolean.TYPE) {
+ return false;
+ }
+ if (returnType == Character.TYPE) {
+ return '\0';
+ }
+ if (returnType == Byte.TYPE) {
+ return (byte) 0;
+ }
+ if (returnType == Short.TYPE) {
+ return (short) 0;
+ }
+ if (returnType == Integer.TYPE) {
+ return 0;
+ }
+ if (returnType == Long.TYPE) {
+ return 0L;
+ }
+ if (returnType == Float.TYPE) {
+ return 0.0F;
+ }
+ return 0.0D;
+ }
+
+ private static class TestSubscriptionProvider extends
AbstractSubscriptionProvider {
+
+ private TestSubscriptionProvider(final int port) {
+ super(
+ new TEndPoint("127.0.0.1", port),
+ "user",
+ "password",
+ null,
+ "consumer",
+ "group",
+ null,
+ null,
+ 1_048_576,
+ 30_000L,
+ CONNECTION_TIMEOUT_MS);
+ }
+
+ @Override
+ protected AbstractSessionBuilder constructSubscriptionSessionBuilder(
+ final String host,
+ final int port,
+ final String username,
+ final String password,
+ final String encryptedPassword,
+ final int thriftMaxFrameSize,
+ final int connectionTimeoutInMs) {
+ final boolean useEncryptedPassword = Objects.nonNull(encryptedPassword);
+ return new SubscriptionTreeSessionBuilder()
+ .host(host)
+ .port(port)
+ .username(username)
+ .password(useEncryptedPassword ? encryptedPassword : password)
+ .useEncryptedPassword(useEncryptedPassword)
+ .thriftMaxFrameSize(thriftMaxFrameSize)
+ .connectionTimeoutInMs(connectionTimeoutInMs);
+ }
+ }
+}
