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 4c2373e2a9c9..061fd948de9b 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 @@ -1410,6 +1410,7 @@ private static void mergeTimeSelectedByTable( private List pollInternal( final Set topicNames, final long timeoutMs) throws SubscriptionException { checkIfFenced(); + boolean syncIfUnavailable = false; providers.acquireReadLock(); try { final AbstractSubscriptionProvider provider = providers.getNextAvailableProvider(); @@ -1417,22 +1418,23 @@ private List pollInternal( 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 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 db0ded818ec2..0beabf8d6145 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 @@ SubscriptionSessionConnection getSessionConnection() throws IoTDBConnectionExcep 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 cae7ded2cf94..565c719149d0 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 @@ 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 subscriptionProviders = new ConcurrentSkipListMap<>(); private final AtomicBoolean isClosing = new AtomicBoolean(false); + private final AtomicBoolean unavailableSyncInFlight = new AtomicBoolean(false); private final Set 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 initialEndpoints; + private long lastUnavailableSyncTimeMs; AbstractSubscriptionProviders(final Set initialEndpoints) { this.initialEndpoints = initialEndpoints; @@ -395,6 +398,43 @@ void sync(final AbstractSubscriptionConsumer consumer) { } } + /** 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/SubscriptionProviderConnectionTimeoutTest.java b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionProviderConnectionTimeoutTest.java new file mode 100644 index 000000000000..ab7a741b07dc --- /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); + } + } +}