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 76796db4782 [Subscription] Bound consumer close when server is
unresponsive (#18659)
76796db4782 is described below
commit 76796db4782c4f555367686f4a77dcf6a6ea976f
Author: Caideyipi <[email protected]>
AuthorDate: Thu Sep 17 11:43:00 2026 +0800
[Subscription] Bound consumer close when server is unresponsive (#18659)
---
.../apache/iotdb/session/SessionConnection.java | 26 ++++
.../iotdb/session/SessionConnectionTest.java | 19 +++
.../SubscriptionSessionConnection.java | 8 ++
.../base/AbstractSubscriptionConsumer.java | 23 ++--
.../base/AbstractSubscriptionProvider.java | 96 ++++++++++---
.../base/AbstractSubscriptionProviders.java | 38 +++++-
.../base/AbstractSubscriptionPullConsumer.java | 86 +++++++-----
.../base/AbstractSubscriptionPushConsumer.java | 20 ++-
.../SubscriptionSessionConnectionTimeoutTest.java | 148 +++++++++++++++++++++
.../base/SubscriptionConsumerLifecycleTest.java | 96 ++++++++++++-
10 files changed, 487 insertions(+), 73 deletions(-)
diff --git
a/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java
b/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java
index 2e144448b58..6bcc41859a8 100644
---
a/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java
+++
b/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java
@@ -30,6 +30,7 @@ import org.apache.iotdb.rpc.RedirectException;
import org.apache.iotdb.rpc.RpcUtils;
import org.apache.iotdb.rpc.StatementExecutionException;
import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.rpc.TimeoutChangeableTransport;
import org.apache.iotdb.rpc.UrlUtils;
import org.apache.iotdb.service.rpc.thrift.IClientRPCService;
import
org.apache.iotdb.service.rpc.thrift.TCreateTimeseriesUsingSchemaTemplateReq;
@@ -326,6 +327,31 @@ public class SessionConnection {
return client;
}
+ protected boolean setTransportTimeout(final int timeoutInMs) {
+ if (!(transport instanceof TimeoutChangeableTransport)) {
+ return false;
+ }
+
+ try {
+ ((TimeoutChangeableTransport) transport).setTimeout(timeoutInMs);
+ return true;
+ } catch (final RuntimeException ignored) {
+ return false;
+ }
+ }
+
+ protected void forceCloseTransport() {
+ if (transport == null) {
+ return;
+ }
+
+ try {
+ transport.close();
+ } catch (final RuntimeException ignored) {
+ // Best effort. The caller must still finish updating its lifecycle
state.
+ }
+ }
+
protected void setTimeZone(String zoneId)
throws StatementExecutionException, IoTDBConnectionException {
final TSStatus status =
diff --git
a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionConnectionTest.java
b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionConnectionTest.java
index fea78466a78..9faa0c6efb6 100644
---
a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionConnectionTest.java
+++
b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionConnectionTest.java
@@ -28,6 +28,7 @@ import org.apache.iotdb.rpc.IoTDBConnectionException;
import org.apache.iotdb.rpc.RedirectException;
import org.apache.iotdb.rpc.StatementExecutionException;
import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.rpc.TimeoutChangeableTFastFramedTransport;
import org.apache.iotdb.service.rpc.thrift.IClientRPCService;
import
org.apache.iotdb.service.rpc.thrift.TCreateTimeseriesUsingSchemaTemplateReq;
import org.apache.iotdb.service.rpc.thrift.TSAppendSchemaTemplateReq;
@@ -205,6 +206,24 @@ public class SessionConnectionTest {
sessionConnection.setStorageGroup("root.test1");
}
+ @Test
+ public void testTransportLifecycleControls() {
+ final TimeoutChangeableTFastFramedTransport timeoutTransport =
+ Mockito.mock(TimeoutChangeableTFastFramedTransport.class);
+ Whitebox.setInternalState(sessionConnection, "transport",
timeoutTransport);
+
+ Assert.assertTrue(sessionConnection.setTransportTimeout(123));
+ Mockito.verify(timeoutTransport).setTimeout(123);
+
+ sessionConnection.forceCloseTransport();
+ Mockito.verify(timeoutTransport).close();
+ }
+
+ @Test
+ public void testSetTransportTimeoutOnUnsupportedTransport() {
+ Assert.assertFalse(sessionConnection.setTransportTimeout(123));
+ }
+
@Test
public void testDeleteStorageGroups()
throws IoTDBConnectionException, StatementExecutionException {
diff --git
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionConnection.java
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionConnection.java
index 3f15f4dcf6a..df3de1331e9 100644
---
a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionConnection.java
+++
b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionConnection.java
@@ -58,4 +58,12 @@ public class SubscriptionSessionConnection extends
SessionConnection {
public TPipeSubscribeResp pipeSubscribe(final TPipeSubscribeReq req) throws
TException {
return client.pipeSubscribe(req);
}
+
+ public boolean setTimeout(final int timeoutInMs) {
+ return setTransportTimeout(timeoutInMs);
+ }
+
+ public void forceClose() {
+ forceCloseTransport();
+ }
}
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 73a51aa9923..8a2c9ff6334 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
@@ -338,24 +338,23 @@ abstract class AbstractSubscriptionConsumer implements
AutoCloseable {
submitEndpointsSyncer();
}
+ protected void prepareClose() {
+ providers.prepareClose();
+ }
+
@Override
- public synchronized void close() {
- if (isClosed.get()) {
+ public void close() {
+ if (!isClosed.compareAndSet(false, true)) {
return;
}
- // close subscription providers
- providers.acquireWriteLock();
- try {
- providers.closeProviders();
- } finally {
- providers.releaseWriteLock();
- }
-
- isClosed.set(true);
-
// mark is released to avoid reopening after closing
isReleased.set(true);
+
+ // Do not wait for the providers write lock here. A poll or heartbeat may
hold it while blocked
+ // in network I/O. prepareClose() bounds or interrupts those RPCs before
providers are detached.
+ providers.prepareClose();
+ providers.closeProviders();
}
boolean isClosed() {
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 cfe741e2d7e..92350097787 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
@@ -54,6 +54,7 @@ import
org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeHeartbeat
import
org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribePollResp;
import
org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeSubscribeResp;
import
org.apache.iotdb.rpc.subscription.payload.response.PipeSubscribeUnsubscribeResp;
+import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeReq;
import org.apache.iotdb.service.rpc.thrift.TPipeSubscribeResp;
import org.apache.iotdb.session.AbstractSessionBuilder;
import org.apache.iotdb.session.subscription.SubscriptionSessionConnection;
@@ -70,11 +71,14 @@ import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.ReentrantLock;
public abstract class AbstractSubscriptionProvider {
private static final Logger LOGGER =
LoggerFactory.getLogger(AbstractSubscriptionProvider.class);
+ private static final int CLOSE_TIMEOUT_IN_MS = 5_000;
+
private static final String STATUS_FORMATTER = "Status code is [%s], status
message is [%s].";
private static final String INTERNAL_ERROR_FORMATTER =
"Internal error occurred. " + STATUS_FORMATTER;
@@ -91,7 +95,9 @@ public abstract class AbstractSubscriptionProvider {
private final Long ownerEpoch;
private final AtomicBoolean isClosed = new AtomicBoolean(true);
+ private final AtomicBoolean isClosing = new AtomicBoolean(false);
private final AtomicBoolean isAvailable = new AtomicBoolean(false);
+ private final ReentrantLock rpcLock = new ReentrantLock();
private final TEndPoint endPoint;
private int dataNodeId;
@@ -149,6 +155,16 @@ public abstract class AbstractSubscriptionProvider {
return session.getSessionConnection();
}
+ private TPipeSubscribeResp pipeSubscribe(final TPipeSubscribeReq req)
+ throws TException, IoTDBConnectionException {
+ rpcLock.lock();
+ try {
+ return getSessionConnection().pipeSubscribe(req);
+ } finally {
+ rpcLock.unlock();
+ }
+ }
+
boolean isAvailable() {
return isAvailable.get();
}
@@ -233,7 +249,7 @@ public abstract class AbstractSubscriptionProvider {
}
final TPipeSubscribeResp resp;
try {
- resp = getSessionConnection().pipeSubscribe(req);
+ resp = pipeSubscribe(req);
} catch (final TException | IoTDBConnectionException e) {
// Assume provider unavailable
LOGGER.warn(
@@ -249,24 +265,72 @@ public abstract class AbstractSubscriptionProvider {
return PipeSubscribeHandshakeResp.fromTPipeSubscribeResp(resp);
}
- synchronized void close() throws SubscriptionException,
IoTDBConnectionException {
- if (isClosed.get()) {
+ void prepareClose() {
+ final SubscriptionSessionConnection connection;
+ try {
+ connection = getSessionConnection();
+ } catch (final IoTDBConnectionException | RuntimeException ignored) {
+ return;
+ }
+ if (connection == null) {
+ return;
+ }
+
+ if (!connection.setTimeout(CLOSE_TIMEOUT_IN_MS)) {
+ setUnavailable();
+ connection.forceClose();
+ return;
+ }
+
+ if (rpcLock.tryLock()) {
+ rpcLock.unlock();
+ } else {
+ // Changing SO_TIMEOUT does not affect a read that is already blocked.
Closing the transport
+ // is the only reliable way to make an in-flight RPC release the
provider immediately.
+ setUnavailable();
+ connection.forceClose();
+ }
+ }
+
+ void close() throws SubscriptionException, IoTDBConnectionException {
+ if (isClosed.get() || !isClosing.compareAndSet(false, true)) {
return;
}
+ SubscriptionSessionConnection connection = null;
try {
- closeInternal(); // throw SubscriptionException
+ connection = getSessionConnection();
+ if (!connection.setTimeout(CLOSE_TIMEOUT_IN_MS) || !rpcLock.tryLock()) {
+ connection.forceClose();
+ return;
+ }
+ try {
+ closeInternal(); // throw SubscriptionException
+ } finally {
+ rpcLock.unlock();
+ }
} finally {
- session.close(); // throw IoTDBConnectionException
- setUnavailable();
- isClosed.set(true);
+ try {
+ if (connection != null) {
+ connection.forceClose();
+ }
+ } finally {
+ try {
+ // The transport is already closed, so Session.close() only updates
local resources and
+ // cannot block on a second closeSession RPC.
+ session.close(); // throw IoTDBConnectionException
+ } finally {
+ setUnavailable();
+ isClosed.set(true);
+ }
+ }
}
}
void closeInternal() throws SubscriptionException {
final TPipeSubscribeResp resp;
try {
- resp =
getSessionConnection().pipeSubscribe(PipeSubscribeCloseReq.toTPipeSubscribeReq());
+ resp = pipeSubscribe(PipeSubscribeCloseReq.toTPipeSubscribeReq());
} catch (final TException | IoTDBConnectionException e) {
// Assume provider unavailable
LOGGER.warn(
@@ -304,7 +368,7 @@ public abstract class AbstractSubscriptionProvider {
final TPipeSubscribeResp resp;
try {
- resp = getSessionConnection().pipeSubscribe(req);
+ resp = pipeSubscribe(req);
} catch (final TException | IoTDBConnectionException e) {
// Assume provider unavailable
LOGGER.warn(
@@ -334,7 +398,7 @@ public abstract class AbstractSubscriptionProvider {
}
final TPipeSubscribeResp resp;
try {
- resp = getSessionConnection().pipeSubscribe(req);
+ resp = pipeSubscribe(req);
} catch (final TException | IoTDBConnectionException e) {
// Assume provider unavailable
LOGGER.warn(
@@ -367,7 +431,7 @@ public abstract class AbstractSubscriptionProvider {
}
final TPipeSubscribeResp resp;
try {
- resp = getSessionConnection().pipeSubscribe(req);
+ resp = pipeSubscribe(req);
} catch (final TException | IoTDBConnectionException e) {
// Assume provider unavailable
LOGGER.warn(
@@ -401,7 +465,7 @@ public abstract class AbstractSubscriptionProvider {
}
final TPipeSubscribeResp resp;
try {
- resp = getSessionConnection().pipeSubscribe(req);
+ resp = pipeSubscribe(req);
} catch (final TException | IoTDBConnectionException e) {
LOGGER.warn(
SubscriptionMessages
@@ -431,7 +495,7 @@ public abstract class AbstractSubscriptionProvider {
}
final TPipeSubscribeResp resp;
try {
- resp = getSessionConnection().pipeSubscribe(req);
+ resp = pipeSubscribe(req);
} catch (final TException | IoTDBConnectionException e) {
LOGGER.warn(
SubscriptionMessages
@@ -461,7 +525,7 @@ public abstract class AbstractSubscriptionProvider {
}
final TPipeSubscribeResp resp;
try {
- resp = getSessionConnection().pipeSubscribe(req);
+ resp = pipeSubscribe(req);
} catch (final TException | IoTDBConnectionException e) {
LOGGER.warn(
SubscriptionMessages
@@ -532,7 +596,7 @@ public abstract class AbstractSubscriptionProvider {
}
final TPipeSubscribeResp resp;
try {
- resp = getSessionConnection().pipeSubscribe(req);
+ resp = pipeSubscribe(req);
} catch (final TException | IoTDBConnectionException e) {
// Assume provider unavailable
LOGGER.warn(
@@ -566,7 +630,7 @@ public abstract class AbstractSubscriptionProvider {
}
final TPipeSubscribeResp resp;
try {
- resp = getSessionConnection().pipeSubscribe(req);
+ resp = pipeSubscribe(req);
} catch (final TException | IoTDBConnectionException e) {
// Assume provider unavailable
LOGGER.warn(
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 d89fe998cdf..0f8e032bec2 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
@@ -38,6 +38,7 @@ import java.util.Objects;
import java.util.Set;
import java.util.SortedMap;
import java.util.concurrent.ConcurrentSkipListMap;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.stream.Collectors;
@@ -47,6 +48,7 @@ final class AbstractSubscriptionProviders {
private final SortedMap<Integer, AbstractSubscriptionProvider>
subscriptionProviders =
new ConcurrentSkipListMap<>();
+ private final AtomicBoolean isClosing = new AtomicBoolean(false);
private int nextDataNodeId = -1;
private final ReentrantReadWriteLock subscriptionProvidersLock = new
ReentrantReadWriteLock(true);
@@ -160,23 +162,51 @@ final class AbstractSubscriptionProviders {
nextDataNodeId = subscriptionProviders.firstKey();
}
- /** Caller should ensure that the method is called in the lock {@link
#acquireWriteLock()}. */
+ /** Detaches and closes the current providers. Terminal consumer close may
call this lock-free. */
void closeProviders() {
- for (final AbstractSubscriptionProvider provider : getAllProviders()) {
+ final List<AbstractSubscriptionProvider> providers = getAllProviders();
+ subscriptionProviders.clear();
+ for (final AbstractSubscriptionProvider provider : providers) {
try {
provider.close();
} catch (final Exception e) {
LOGGER.warn(SubscriptionMessages.PROVIDER_CLOSE_FAILED, provider, e,
e);
}
}
- subscriptionProviders.clear();
+ }
+
+ void prepareClose() {
+ isClosing.set(true);
+ for (final AbstractSubscriptionProvider provider : getAllProviders()) {
+ provider.prepareClose();
+ }
}
/** Caller should ensure that the method is called in the lock {@link
#acquireWriteLock()}. */
void addProvider(final int dataNodeId, final AbstractSubscriptionProvider
provider) {
+ if (isClosing.get()) {
+ closeProviderAddedDuringClosing(provider);
+ return;
+ }
+
+ subscriptionProviders.put(dataNodeId, provider);
+ if (isClosing.get()) {
+ subscriptionProviders.remove(dataNodeId, provider);
+ closeProviderAddedDuringClosing(provider);
+ return;
+ }
+
// the subscription provider is opened
LOGGER.info(SubscriptionMessages.ADD_NEW_PROVIDER, provider);
- subscriptionProviders.put(dataNodeId, provider);
+ }
+
+ private void closeProviderAddedDuringClosing(final
AbstractSubscriptionProvider provider) {
+ provider.prepareClose();
+ try {
+ provider.close();
+ } catch (final Exception e) {
+ LOGGER.warn(SubscriptionMessages.PROVIDER_CLOSE_FAILED, provider, e, e);
+ }
}
/** Caller should ensure that the method is called in the lock {@link
#acquireWriteLock()}. */
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 811b4228c1e..61d558cd829 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
@@ -87,6 +87,7 @@ public abstract class AbstractSubscriptionPullConsumer
extends AbstractSubscript
private final EmptyPollLogThrottler emptyPollLogThrottler = new
EmptyPollLogThrottler();
private final AtomicBoolean isClosed = new AtomicBoolean(true);
+ private final AtomicBoolean isClosing = new AtomicBoolean(false);
@Override
boolean isClosed() {
@@ -154,57 +155,70 @@ public abstract class AbstractSubscriptionPullConsumer
extends AbstractSubscript
}
@Override
- public synchronized void close() {
- if (isClosed.get()) {
+ public void close() {
+ if (isClosed.get() || !isClosing.compareAndSet(false, true)) {
return;
}
- if (!processors.isEmpty()) {
- if (autoCommit) {
- final List<SubscriptionMessage> drainedMessages =
drainProcessorPipeline();
- if (!drainedMessages.isEmpty()) {
+ try {
+ synchronized (this) {
+ if (isClosed.get()) {
+ return;
+ }
+
+ List<SubscriptionMessage> drainedProcessorMessages =
Collections.emptyList();
+ if (!processors.isEmpty()) {
+ drainedProcessorMessages = drainProcessorPipeline();
+ if (!autoCommit && !drainedProcessorMessages.isEmpty()) {
+ pendingDrainedMessages.addAll(drainedProcessorMessages);
+ }
+ }
+
+ // In manual-commit mode, preserve the existing retry contract:
validate before making the
+ // terminal close state visible so the caller can drain and commit,
then retry close().
+ if (!autoCommit) {
+ ensureNoManualBufferedMessagesOnClose();
+ }
+
+ // Publish the terminal state before touching the network. Workers
stop starting new RPCs,
+ // while prepareClose() bounds subsequent requests and interrupts any
read already blocked.
+ isClosed.set(true);
+ prepareClose();
+
+ if (autoCommit && !drainedProcessorMessages.isEmpty()) {
try {
- commitSync(drainedMessages);
+ commitSync(drainedProcessorMessages);
} catch (final SubscriptionException e) {
LOGGER.warn(
SubscriptionMessages.LOG_FAILED_COMMIT_DRAINED_PROCESSOR_MESSAGES_CLOSE_4264DB35,
e);
}
}
- } else {
- final List<SubscriptionMessage> drainedMessages =
drainProcessorPipeline();
- if (!drainedMessages.isEmpty()) {
- pendingDrainedMessages.addAll(drainedMessages);
- }
- ensureNoManualBufferedMessagesOnClose();
- }
- }
- if (autoCommit && !pendingDrainedMessages.isEmpty()) {
- final List<SubscriptionMessage> drainedMessages =
drainPendingDrainedMessages();
- if (!drainedMessages.isEmpty()) {
- try {
- commitSync(drainedMessages);
- } catch (final SubscriptionException e) {
- LOGGER.warn(
- SubscriptionMessages
-
.LOG_FAILED_COMMIT_PENDING_DRAINED_PROCESSOR_MESSAGES_CLOSE_644B5DDD,
- e);
+ if (autoCommit && !pendingDrainedMessages.isEmpty()) {
+ final List<SubscriptionMessage> drainedMessages =
drainPendingDrainedMessages();
+ if (!drainedMessages.isEmpty()) {
+ try {
+ commitSync(drainedMessages);
+ } catch (final SubscriptionException e) {
+ LOGGER.warn(
+ SubscriptionMessages
+
.LOG_FAILED_COMMIT_PENDING_DRAINED_PROCESSOR_MESSAGES_CLOSE_644B5DDD,
+ e);
+ }
+ }
}
- }
- }
- if (!autoCommit) {
- ensureNoManualBufferedMessagesOnClose();
- }
+ if (autoCommit) {
+ // commit all uncommitted messages
+ commitAllUncommittedMessages();
+ }
- if (autoCommit) {
- // commit all uncommitted messages
- commitAllUncommittedMessages();
+ super.close();
+ }
+ } finally {
+ isClosing.set(false);
}
-
- isClosed.set(true);
- super.close();
}
/////////////////////////////// poll & commit ///////////////////////////////
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 f3c7d38c4dc..a06a2e9cc4e 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
@@ -66,6 +66,7 @@ public abstract class AbstractSubscriptionPushConsumer
extends AbstractSubscript
private final EmptyPollLogThrottler emptyPollLogThrottler = new
EmptyPollLogThrottler();
private final AtomicBoolean isClosed = new AtomicBoolean(true);
+ private final AtomicBoolean isClosing = new AtomicBoolean(false);
protected AbstractSubscriptionPushConsumer(
final AbstractSubscriptionPushConsumerBuilder builder) {
@@ -142,13 +143,24 @@ public abstract class AbstractSubscriptionPushConsumer
extends AbstractSubscript
}
@Override
- public synchronized void close() {
- if (isClosed.get()) {
+ public void close() {
+ if (isClosed.get() || !isClosing.compareAndSet(false, true)) {
return;
}
- isClosed.set(true);
- super.close();
+ try {
+ synchronized (this) {
+ if (isClosed.get()) {
+ return;
+ }
+
+ isClosed.set(true);
+ prepareClose();
+ super.close();
+ }
+ } finally {
+ isClosing.set(false);
+ }
}
@Override
diff --git
a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/SubscriptionSessionConnectionTimeoutTest.java
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/SubscriptionSessionConnectionTimeoutTest.java
new file mode 100644
index 00000000000..8e0446cd71c
--- /dev/null
+++
b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/SubscriptionSessionConnectionTimeoutTest.java
@@ -0,0 +1,148 @@
+/*
+ * 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;
+
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.rpc.TElasticFramedTransport;
+import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeCloseReq;
+import org.apache.iotdb.service.rpc.thrift.IClientRPCService;
+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.thrift.TException;
+import org.apache.thrift.server.TSimpleServer;
+import org.apache.thrift.transport.TServerSocket;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.lang.reflect.Proxy;
+import java.net.InetAddress;
+import java.net.ServerSocket;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+public class SubscriptionSessionConnectionTimeoutTest {
+
+ private static final int READ_TIMEOUT_MS = 200;
+
+ @Test(timeout = 10_000)
+ public void testPipeSubscribeReadTimeout() throws Exception {
+ final CountDownLatch pipeSubscribeStarted = new CountDownLatch(1);
+ final CountDownLatch releasePipeSubscribe = new CountDownLatch(1);
+ 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 "pipeSubscribe":
+ pipeSubscribeStarted.countDown();
+ releasePipeSubscribe.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 TElasticFramedTransport.Factory()));
+ final Thread serverThread = new Thread(server::serve,
"subscription-timeout-test-server");
+ serverThread.setDaemon(true);
+ serverThread.start();
+
+ final SubscriptionSessionWrapper session =
+ new SubscriptionSessionWrapper(
+ new SubscriptionTreeSessionBuilder()
+ .host("127.0.0.1")
+ .port(socket.getLocalPort())
+ .connectionTimeoutInMs(0));
+ SubscriptionSessionConnection connection = null;
+ try {
+ session.open();
+ connection = session.getSessionConnection();
+ Assert.assertTrue(connection.setTimeout(READ_TIMEOUT_MS));
+
+ final long startNanos = System.nanoTime();
+ try {
+ connection.pipeSubscribe(PipeSubscribeCloseReq.toTPipeSubscribeReq());
+ Assert.fail("Expected the unresponsive pipeSubscribe RPC to time out");
+ } catch (final TException expected) {
+ final long elapsedMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime()
- startNanos);
+ Assert.assertTrue("RPC did not honor the read timeout: " + elapsedMs,
elapsedMs < 3_000);
+ }
+ Assert.assertTrue(pipeSubscribeStarted.await(1, TimeUnit.SECONDS));
+ } finally {
+ if (connection != null) {
+ connection.forceClose();
+ }
+ session.close();
+ releasePipeSubscribe.countDown();
+ server.stop();
+ serverThread.join(2_000);
+ }
+ }
+
+ 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;
+ }
+}
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 68f714f2ce5..df42e2b244c 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
@@ -32,6 +32,11 @@ import org.junit.Test;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
+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.function.BooleanSupplier;
public class SubscriptionConsumerLifecycleTest {
@@ -92,6 +97,34 @@ public class SubscriptionConsumerLifecycleTest {
Assert.assertTrue(consumer.closedStatesDuringClose.get(0));
}
+ @Test
+ public void testConcurrentPullConsumerCloseReturnsWithoutWaiting() throws
Exception {
+ final CountDownLatch providerCloseStarted = new CountDownLatch(1);
+ final CountDownLatch allowProviderClose = new CountDownLatch(1);
+ final TestPullConsumer consumer =
+ new TestPullConsumer(providerCloseStarted, allowProviderClose);
+ final ExecutorService executor = Executors.newFixedThreadPool(2);
+
+ try {
+ consumer.open();
+
+ final Future<?> firstClose = executor.submit(consumer::close);
+ Assert.assertTrue(providerCloseStarted.await(5, TimeUnit.SECONDS));
+
+ final Future<?> concurrentClose = executor.submit(consumer::close);
+ concurrentClose.get(1, TimeUnit.SECONDS);
+ Assert.assertFalse(firstClose.isDone());
+
+ allowProviderClose.countDown();
+ firstClose.get(5, TimeUnit.SECONDS);
+ Assert.assertEquals(1, consumer.closedStatesDuringClose.size());
+ } finally {
+ allowProviderClose.countDown();
+ executor.shutdownNow();
+ executor.awaitTermination(5, TimeUnit.SECONDS);
+ }
+ }
+
private static class TestPushConsumer extends
AbstractSubscriptionPushConsumer {
private final List<Boolean> closedStatesDuringHandshake = new
ArrayList<>();
@@ -144,8 +177,15 @@ public class SubscriptionConsumerLifecycleTest {
private final List<Boolean> closedStatesDuringHandshake = new
ArrayList<>();
private final List<Boolean> closedStatesDuringClose = new ArrayList<>();
+ private final CountDownLatch providerCloseStarted;
+ private final CountDownLatch allowProviderClose;
private TestPullConsumer() {
+ this(null, null);
+ }
+
+ private TestPullConsumer(
+ final CountDownLatch providerCloseStarted, final CountDownLatch
allowProviderClose) {
super(
new AbstractSubscriptionPullConsumerBuilder()
.host(HOST)
@@ -155,6 +195,8 @@ public class SubscriptionConsumerLifecycleTest {
.heartbeatIntervalMs(LONG_INTERVAL_MS)
.endpointsSyncIntervalMs(LONG_INTERVAL_MS)
.autoCommit(false));
+ this.providerCloseStarted = providerCloseStarted;
+ this.allowProviderClose = allowProviderClose;
}
@Override
@@ -184,7 +226,9 @@ public class SubscriptionConsumerLifecycleTest {
connectionTimeoutInMs,
this::isClosed,
closedStatesDuringHandshake,
- closedStatesDuringClose);
+ closedStatesDuringClose,
+ providerCloseStarted,
+ allowProviderClose);
}
}
@@ -193,6 +237,8 @@ public class SubscriptionConsumerLifecycleTest {
private final BooleanSupplier consumerClosedSupplier;
private final List<Boolean> closedStatesDuringHandshake;
private final List<Boolean> closedStatesDuringClose;
+ private final CountDownLatch providerCloseStarted;
+ private final CountDownLatch allowProviderClose;
private TestSubscriptionProvider(
final TEndPoint endPoint,
@@ -209,6 +255,42 @@ public class SubscriptionConsumerLifecycleTest {
final BooleanSupplier consumerClosedSupplier,
final List<Boolean> closedStatesDuringHandshake,
final List<Boolean> closedStatesDuringClose) {
+ this(
+ endPoint,
+ username,
+ password,
+ encryptedPassword,
+ consumerId,
+ consumerGroupId,
+ ownerId,
+ ownerEpoch,
+ thriftMaxFrameSize,
+ heartbeatIntervalMs,
+ connectionTimeoutInMs,
+ consumerClosedSupplier,
+ closedStatesDuringHandshake,
+ closedStatesDuringClose,
+ null,
+ null);
+ }
+
+ private TestSubscriptionProvider(
+ final TEndPoint endPoint,
+ final String username,
+ final String password,
+ final String encryptedPassword,
+ final String consumerId,
+ final String consumerGroupId,
+ final String ownerId,
+ final Long ownerEpoch,
+ final int thriftMaxFrameSize,
+ final long heartbeatIntervalMs,
+ final int connectionTimeoutInMs,
+ final BooleanSupplier consumerClosedSupplier,
+ final List<Boolean> closedStatesDuringHandshake,
+ final List<Boolean> closedStatesDuringClose,
+ final CountDownLatch providerCloseStarted,
+ final CountDownLatch allowProviderClose) {
super(
endPoint,
username,
@@ -224,6 +306,8 @@ public class SubscriptionConsumerLifecycleTest {
this.consumerClosedSupplier = consumerClosedSupplier;
this.closedStatesDuringHandshake = closedStatesDuringHandshake;
this.closedStatesDuringClose = closedStatesDuringClose;
+ this.providerCloseStarted = providerCloseStarted;
+ this.allowProviderClose = allowProviderClose;
}
@Override
@@ -255,6 +339,16 @@ public class SubscriptionConsumerLifecycleTest {
@Override
synchronized void close() {
closedStatesDuringClose.add(consumerClosedSupplier.getAsBoolean());
+ if (Objects.nonNull(providerCloseStarted)) {
+ providerCloseStarted.countDown();
+ }
+ if (Objects.nonNull(allowProviderClose)) {
+ try {
+ allowProviderClose.await();
+ } catch (final InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ }
setUnavailable();
}
