diff --git a/integration-test/src/test/java/org/apache/iotdb/session/it/SessionSeedNodeFallbackIT.java b/integration-test/src/test/java/org/apache/iotdb/session/it/SessionSeedNodeFallbackIT.java new file mode 100644 index 0000000000000..81df4e4fcf114 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/session/it/SessionSeedNodeFallbackIT.java @@ -0,0 +1,244 @@ +/* + * 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.it; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.cluster.NodeStatus; +import org.apache.iotdb.consensus.ConsensusFactory; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.isession.pool.SessionDataSetWrapper; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.env.cluster.EnvUtils; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.itbase.category.ClusterIT; +import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.rpc.UrlUtils; +import org.apache.iotdb.session.NodesSupplier; +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.pool.SessionPool; + +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.lang.reflect.Field; +import java.time.ZoneId; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +@RunWith(IoTDBTestRunner.class) +@Category(ClusterIT.class) +public class SessionSeedNodeFallbackIT { + @Before + public void setUp() { + EnvFactory.getEnv() + .getConfig() + .getCommonConfig() + .setSchemaReplicationFactor(3) + .setDataReplicationFactor(3) + .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS) + .setDataRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS); + EnvFactory.getEnv().initClusterEnvironment(1, 3); + } + + @After + public void tearDown() { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + @Test + public void testQueriesDuringRollingRestart() throws Exception { + SessionPool pool = + new SessionPool.Builder().nodeUrls(urls()).maxSize(1).connectionTimeoutInMs(1000).build(); + Session session = new Session.Builder().nodeUrls(urls()).build(); + try { + // Initial Ratis region creation can exceed one second; do not retry its DDL prematurely. + session.open(false, 30_000); + session.executeNonQueryStatement("CREATE DATABASE root.seed_fallback"); + session.executeNonQueryStatement( + "CREATE TIMESERIES root.seed_fallback.d.s WITH DATATYPE=INT32, ENCODING=RLE"); + session.executeNonQueryStatement("INSERT INTO root.seed_fallback.d(time,s) VALUES(1,42)"); + assertQuery(session, pool); + NodesSupplier sessionSupplier = supplier(session); + NodesSupplier poolSupplier = supplier(pool); + for (int index = 0; index < 3; index++) { + EnvFactory.getEnv().shutdownDataNode(index); + awaitStatus(index, EnvUtils.getNodeStatusAfterLocalStop()); + sessionSupplier.run(); + poolSupplier.run(); + assertFalse(sessionSupplier.get().contains(endpoint(index))); + assertFalse(poolSupplier.get().contains(endpoint(index))); + assertQuery(session, pool); + + EnvFactory.getEnv().startDataNode(index); + awaitStatus(index, NodeStatus.Running); + sessionSupplier.run(); + poolSupplier.run(); + assertTrue(sessionSupplier.get().contains(endpoint(index))); + assertTrue(poolSupplier.get().contains(endpoint(index))); + assertQuery(session, pool); + } + } finally { + try { + session.close(); + } catch (IoTDBConnectionException e) { + // Cached connections may still belong to a stopped process. Only cleanup tolerates this; + // all query and discovery assertions above must succeed without ignoring failures. + } finally { + pool.close(); + } + } + } + + @Test + public void testBusinessAndBackgroundRecoveryThroughOmittedSeed() throws Exception { + EnvFactory.getEnv().shutdownDataNode(0); + awaitStatus(0, EnvUtils.getNodeStatusAfterLocalStop()); + // Control refresh timing so neither supplier learns that node 0 recovered before failover. + ScheduledThreadPoolExecutor executor = + new ScheduledThreadPoolExecutor(1) { + @Override + public ScheduledFuture scheduleAtFixedRate( + Runnable command, long initialDelay, long period, TimeUnit unit) { + return super.scheduleAtFixedRate(command, 1, 1, TimeUnit.DAYS); + } + }; + NodesSupplier business = createSupplier(executor); + NodesSupplier background = createSupplier(executor); + try (Session session = new Session.Builder().nodeUrls(urls()).enableAutoFetch(false).build()) { + business.run(); + background.run(); + assertEquals( + new HashSet<>(Arrays.asList(endpoint(1), endpoint(2))), new HashSet<>(business.get())); + assertEquals(business.get(), background.get()); + session.open(false, 1000, new ConcurrentHashMap<>(), business); + + EnvFactory.getEnv().startDataNode(0); + awaitStatus(0, NodeStatus.Running); + EnvFactory.getEnv().getDataNodeWrapper(1).stopForcibly(); + awaitStatus(1, NodeStatus.Unknown); + EnvFactory.getEnv().shutdownDataNode(2); + awaitStatus(2, EnvUtils.getNodeStatusAfterLocalStop()); + assertFalse(business.get().contains(endpoint(0))); + assertFalse(background.get().contains(endpoint(0))); + + // Both discovered endpoints are dead. Only the initial seed omitted by discovery is live. + try (SessionDataSet result = session.executeQueryStatement("SHOW AVAILABLE URLS")) { + assertEquals(Collections.singleton(endpoint(0)), readEndpoints(result)); + } + background.run(); + assertEquals(Collections.singletonList(endpoint(0)), background.get()); + assertEquals(3, background.getSeedNodes().size()); + // Wait for the immediate refresh requested by business reconnection, not a 60-second poll. + executor.submit(() -> {}).get(10, TimeUnit.SECONDS); + assertEquals(Collections.singletonList(endpoint(0)), business.get()); + } finally { + business.close(); + background.close(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + private void assertQuery(Session session, SessionPool pool) throws Exception { + String sql = "SELECT s FROM root.seed_fallback.d"; + try (SessionDataSet result = session.executeQueryStatement(sql)) { + assertTrue(result.hasNext()); + assertEquals(42, result.next().getFields().get(0).getIntV()); + assertFalse(result.hasNext()); + } + try (SessionDataSetWrapper result = pool.executeQueryStatement(sql)) { + assertTrue(result.hasNext()); + assertEquals(42, result.next().getFields().get(0).getIntV()); + assertFalse(result.hasNext()); + } + } + + private NodesSupplier supplier(Object owner) throws Exception { + Field field = owner.getClass().getDeclaredField("availableNodes"); + field.setAccessible(true); + return (NodesSupplier) field.get(owner); + } + + private NodesSupplier createSupplier(ScheduledThreadPoolExecutor executor) { + List seeds = Arrays.asList(endpoint(0), endpoint(1), endpoint(2)); + return NodesSupplier.createNodeSupplier( + seeds, + executor, + "root", + "root", + ZoneId.of("UTC"), + SessionConfig.DEFAULT_INITIAL_BUFFER_CAPACITY, + SessionConfig.DEFAULT_MAX_FRAME_SIZE, + 1000, + false, + null, + null, + null, + null, + "TLS", + false, + "V_1_0"); + } + + private Set readEndpoints(SessionDataSet dataSet) throws Exception { + Set result = new HashSet<>(); + SessionDataSet.DataIterator iterator = dataSet.iterator(); + while (iterator.next()) { + result.add(new TEndPoint(iterator.getString("RpcAddress"), iterator.getInt("RpcPort"))); + } + return result; + } + + private TEndPoint endpoint(int index) { + return new TEndPoint( + EnvFactory.getEnv().getDataNodeWrapper(index).getIp(), + EnvFactory.getEnv().getDataNodeWrapper(index).getPort()); + } + + private List urls() { + return Arrays.asList(endpoint(0), endpoint(1), endpoint(2)).stream() + .map(UrlUtils::convertTEndPointIpv4AndIpv6Url) + .collect(Collectors.toCollection(ArrayList::new)); + } + + private void awaitStatus(int index, NodeStatus status) { + EnvFactory.getEnv() + .ensureNodeStatus( + Collections.singletonList(EnvFactory.getEnv().getDataNodeWrapper(index)), + Collections.singletonList(status)); + } +} diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/INodeSupplier.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/INodeSupplier.java index 60a2a5adb4347..9c8d32f946120 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/INodeSupplier.java +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/INodeSupplier.java @@ -21,6 +21,7 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import java.util.Collections; import java.util.List; import java.util.Optional; import java.util.function.Supplier; @@ -30,4 +31,14 @@ public interface INodeSupplier extends Supplier> { void close(); Optional getQueryEndPoint(); + + /** Initial addresses to try after all discovered endpoints have failed. */ + default List getSeedNodes() { + return Collections.emptyList(); + } + + /** Request a discovery refresh after reconnecting through a seed. */ + default void requestRefresh() { + // Suppliers without automatic discovery have nothing to refresh. + } } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/NodesSupplier.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/NodesSupplier.java index 81d83450f72e6..3a52c861938f9 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/NodesSupplier.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/NodesSupplier.java @@ -30,12 +30,17 @@ import java.time.ZoneId; import java.util.ArrayList; +import java.util.Collections; import java.util.HashSet; +import java.util.LinkedHashSet; import java.util.List; import java.util.Optional; -import java.util.concurrent.CopyOnWriteArrayList; +import java.util.Set; +import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.stream.Collectors; public class NodesSupplier implements INodeSupplier, Runnable { @@ -54,10 +59,13 @@ public class NodesSupplier implements INodeSupplier, Runnable { private static final int FETCH_SIZE = 10_000; - // availableNodes won't be updated frequently, so we use CopyOnWriteArrayList which is thread-safe - // and is optimized for scenarios of reading more and writing less + // Publish complete snapshots; discovery must never overwrite the initial recovery addresses. @SuppressWarnings("java:S3077") - private volatile List availableNodes = new CopyOnWriteArrayList<>(); + private volatile List availableNodes; + + private final List seedNodes; + private final ScheduledExecutorService executorService; + private final AtomicBoolean refreshRequested = new AtomicBoolean(); private final boolean useSSL; private final String trustStore; @@ -108,6 +116,7 @@ public static NodesSupplier createNodeSupplier( NodesSupplier nodesSupplier = new NodesSupplier( endPointList, + executorService, userName, password, zoneId, @@ -133,6 +142,7 @@ public static NodesSupplier createNodeSupplier( private NodesSupplier( List endPointList, + ScheduledExecutorService executorService, String userName, String password, ZoneId zoneId, @@ -147,7 +157,11 @@ private NodesSupplier( String sslProtocol, boolean enableRPCCompression, String version) { - this.availableNodes.addAll(new HashSet<>(endPointList)); + this.seedNodes = + Collections.unmodifiableList( + endPointList.stream().distinct().map(TEndPoint::deepCopy).collect(Collectors.toList())); + this.availableNodes = seedNodes; + this.executorService = executorService; this.userName = userName; this.password = password; this.useSSL = useSSL; @@ -172,24 +186,54 @@ public List get() { } @Override - public void run() { + public List getSeedNodes() { + return seedNodes; + } + + @Override + public void requestRefresh() { + if (!closed && refreshRequested.compareAndSet(false, true)) { + try { + executorService.execute( + () -> { + try { + run(); + } finally { + refreshRequested.set(false); + } + }); + } catch (RejectedExecutionException e) { + // The session or pool may have shut down its executor concurrently. + refreshRequested.set(false); + } + } + } + + @Override + public synchronized void run() { if (closed) { if (client != null) { destroyCurrentClient(); } return; } - if (client == null) { - for (TEndPoint endPoint : availableNodes) { - if (createConnection(endPoint)) { - break; - } + Set attempted = new HashSet<>(); + if (client != null) { + attempted.add(client.endPoint); + if (updateDataNodeList()) { + return; } + destroyCurrentClient(); } - // make sure the following code block can run thread-safely with close() method. - synchronized (this) { - if (client != null && !updateDataNodeList()) { + // A stale discovery snapshot can omit the only node that has finished a rolling restart. + Set candidates = new LinkedHashSet<>(availableNodes); + candidates.addAll(seedNodes); + for (TEndPoint endPoint : candidates) { + if (attempted.add(endPoint)) { + if (createConnection(endPoint) && updateDataNodeList()) { + return; + } destroyCurrentClient(); } } @@ -235,10 +279,11 @@ public void close() { @Override public Optional getQueryEndPoint() { - if (availableNodes == null || availableNodes.isEmpty()) { + List snapshot = availableNodes; + if (snapshot.isEmpty()) { return Optional.empty(); } else { - return Optional.of(policy.chooseOne(get())); + return Optional.of(policy.chooseOne(snapshot)); } } @@ -266,9 +311,7 @@ private void updateAvailableNodes(SessionDataSet sessionDataSet) throws Exceptio } } } - // replace the older ones. - if (!res.isEmpty()) { - availableNodes = res; - } + // An empty successful result is authoritative too; seeds remain available for recovery. + availableNodes = Collections.unmodifiableList(res); } } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java index a6deda174913d..b9e8fe3652ecb 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java @@ -125,6 +125,7 @@ public class Session implements ISession { new LinkedBlockingQueue<>(SessionConfig.DEFAULT_SESSION_EXECUTOR_TASK_NUM), ThreadUtils.createThreadFactory("SessionExecutor", true)); protected List nodeUrls; + private final List initialNodeUrls; protected String username; protected String password; protected boolean useEncryptedPassword; @@ -367,6 +368,8 @@ public Session( boolean enableRedirection, Version version) { this.defaultEndPoint = new TEndPoint(host, rpcPort); + this.initialNodeUrls = + Collections.singletonList(UrlUtils.convertTEndPointIpv4AndIpv6Url(defaultEndPoint)); this.username = username; this.password = password; this.fetchSize = fetchSize; @@ -434,8 +437,9 @@ public Session( if (nodeUrls.isEmpty()) { throw new IllegalArgumentException(SessionMessages.NODE_URLS_EMPTY); } - nodeUrls = shuffleNodeUrls(nodeUrls); + nodeUrls = shuffleNodeUrls(new ArrayList<>(nodeUrls)); this.nodeUrls = nodeUrls; + this.initialNodeUrls = Collections.unmodifiableList(new ArrayList<>(nodeUrls)); this.username = username; this.password = password; this.fetchSize = fetchSize; @@ -451,11 +455,13 @@ public Session(AbstractSessionBuilder builder) { if (builder.nodeUrls.isEmpty()) { throw new IllegalArgumentException(SessionMessages.NODE_URLS_EMPTY); } - builder.nodeUrls = shuffleNodeUrls(builder.nodeUrls); - this.nodeUrls = builder.nodeUrls; + this.nodeUrls = shuffleNodeUrls(new ArrayList<>(builder.nodeUrls)); + this.initialNodeUrls = Collections.unmodifiableList(new ArrayList<>(nodeUrls)); this.enableQueryRedirection = true; } else { this.defaultEndPoint = new TEndPoint(builder.host, builder.rpcPort); + this.initialNodeUrls = + Collections.singletonList(UrlUtils.convertTEndPointIpv4AndIpv6Url(defaultEndPoint)); this.enableQueryRedirection = builder.enableRedirection; } this.enableThriftRpcCompaction = builder.isThriftRpcCompactionEnabled; @@ -599,11 +605,7 @@ private static List shuffleNodeUrls(List endPoints) { } private List getNodeUrls() { - if (defaultEndPoint != null) { - return Collections.singletonList(defaultEndPoint); - } else { - return SessionUtils.parseSeedNodeUrls(nodeUrls); - } + return SessionUtils.parseSeedNodeUrls(initialNodeUrls); } @Override 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 7b6ec3b03dfea..1dde4a5c35ec7 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 @@ -22,6 +22,7 @@ import org.apache.iotdb.common.rpc.thrift.TAggregationType; import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.isession.SessionDataSet; import org.apache.iotdb.rpc.DeepCopyRpcTransportFactory; @@ -84,6 +85,7 @@ import java.time.ZoneId; import java.util.ArrayList; import java.util.Collections; +import java.util.LinkedHashSet; import java.util.List; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; @@ -1143,18 +1145,18 @@ private boolean reconnect() { for (int i = 1; i <= SessionConfig.RETRY_NUM; i++) { if (transport != null) { transport.close(); - endPointList = availableNodes.get(); - int currHostIndex = random.nextInt(endPointList.size()); - int tryHostNum = 0; - for (int j = currHostIndex; j < endPointList.size(); j++) { - if (tryHostNum == endPointList.size()) { - break; - } - this.endPoint = endPointList.get(j); - if (j == endPointList.size() - 1) { - j = -1; - } - tryHostNum++; + List discoveredNodes = availableNodes.get(); + endPointList = new ArrayList<>(discoveredNodes); + if (!endPointList.isEmpty()) { + Collections.rotate(endPointList, -random.nextInt(endPointList.size())); + } + if (availableNodes instanceof INodeSupplier) { + endPointList.addAll(((INodeSupplier) availableNodes).getSeedNodes()); + } + // Try every discovered endpoint before falling back to seeds, once per endpoint per round. + endPointList = new ArrayList<>(new LinkedHashSet<>(endPointList)); + for (TEndPoint candidate : endPointList) { + this.endPoint = candidate; try { init( endPoint, @@ -1165,6 +1167,9 @@ private boolean reconnect() { session.keyStorePwd, session.sslProtocol); connectedSuccess = true; + if (!discoveredNodes.contains(endPoint) && availableNodes instanceof INodeSupplier) { + ((INodeSupplier) availableNodes).requestRefresh(); + } } catch (IoTDBConnectionException e) { logger.warn(SessionMessages.NODE_DOWN_TRY_NEXT, endPoint); continue; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java index 32feac2cea5e1..701979c47ac6b 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java @@ -52,6 +52,7 @@ import java.io.IOException; import java.time.ZoneId; +import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Map; @@ -564,7 +565,7 @@ public SessionPool(AbstractSessionPoolBuilder builder) { if (builder.nodeUrls.isEmpty()) { throw new IllegalArgumentException(SessionMessages.NODE_URLS_EMPTY); } - this.nodeUrls = builder.nodeUrls; + this.nodeUrls = Collections.unmodifiableList(new ArrayList<>(builder.nodeUrls)); this.host = null; this.port = -1; this.formattedNodeUrls = builder.nodeUrls.toString(); diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionSeedNodeFallbackTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionSeedNodeFallbackTest.java new file mode 100644 index 0000000000000..bb1d3bfd5eec4 --- /dev/null +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionSeedNodeFallbackTest.java @@ -0,0 +1,411 @@ +/* + * 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; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.rpc.DeepCopyRpcTransportFactory; +import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.iotdb.rpc.UrlUtils; +import org.apache.iotdb.service.rpc.thrift.IClientRPCService; +import org.apache.iotdb.service.rpc.thrift.TSExecuteStatementResp; +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.protocol.TBinaryProtocol; +import org.apache.thrift.server.TThreadPoolServer; +import org.apache.thrift.transport.TServerSocket; +import org.apache.thrift.transport.TTransport; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.read.common.block.TsBlockBuilder; +import org.apache.tsfile.read.common.block.column.TsBlockSerde; +import org.apache.tsfile.utils.BytesUtils; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.powermock.reflect.Whitebox; + +import java.net.ServerSocket; +import java.time.ZoneId; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +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.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Exercise discovery and reconnection over real Thrift connections with controlled availability. + */ +public class SessionSeedNodeFallbackTest { + private final List nodes = new ArrayList<>(); + private final List attempts = new CopyOnWriteArrayList<>(); + private final List refreshes = new ArrayList<>(); + private final List connections = new ArrayList<>(); + private volatile List advertised; + private ScheduledExecutorService executor; + private NodesSupplier supplier; + + @Before + public void setUp() throws Exception { + for (int i = 0; i < 3; i++) { + nodes.add(new TestNode()); + } + advertised = endpoints(0, 1, 2); + executor = mock(ScheduledExecutorService.class); + doAnswer( + invocation -> { + refreshes.add(invocation.getArgument(0)); + return null; + }) + .when(executor) + .execute(any(Runnable.class)); + supplier = createSupplier(endpoints(0, 1, 2)); + } + + @After + public void tearDown() throws Exception { + for (SessionConnection connection : connections) { + connection.forceCloseTransport(); + } + if (supplier != null) { + supplier.close(); + } + for (TestNode node : nodes) { + node.close(); + } + } + + @Test + public void testBackgroundRefreshRecoversFromStaleSnapshot() { + nodes.get(0).online = false; + advertised = endpoints(1, 2); + supplier.run(); + assertEquals(advertised, supplier.get()); + + // Node 1 has finished upgrading, but both cached candidates are now unavailable. + nodes.get(0).online = true; + nodes.get(1).online = false; + nodes.get(2).online = false; + advertised = endpoints(0); + attempts.clear(); + supplier.run(); + assertEquals(endpoints(2, 0), attempts); + assertEquals(endpoints(0), supplier.get()); + assertEquals(endpoints(0, 1, 2), supplier.getSeedNodes()); + } + + @Test + public void testBusinessReconnectFallsBackAndRequestsRefresh() throws Exception { + advertised = endpoints(1, 2); + supplier.run(); + nodes.get(1).online = false; + nodes.get(2).online = false; + advertised = endpoints(0); + attempts.clear(); + SessionConnection connection = newConnection(); + assertTrue(reconnect(connection)); + assertEquals(nodes.get(0).endpoint, connection.getEndPoint()); + assertEquals(3, attempts.size()); + assertEquals(nodes.get(0).endpoint, attempts.get(2)); + assertEquals(1, refreshes.size()); + refreshes.remove(0).run(); + assertEquals(endpoints(0), supplier.get()); + } + + @Test + public void testDiscoveredNodeIsPreferredToHealthySeed() throws Exception { + advertised = endpoints(1); + supplier.run(); + attempts.clear(); + assertTrue(reconnect(newConnection())); + assertEquals(endpoints(1), attempts); + assertTrue(refreshes.isEmpty()); + for (int i = 0; i < 6; i++) { + assertEquals(nodes.get(1).endpoint, supplier.getQueryEndPoint().get()); + } + } + + @Test + public void testEmptyDiscoveryFallsBackWithoutLosingSeeds() throws Exception { + advertised = Collections.emptyList(); + supplier.run(); + assertTrue(supplier.get().isEmpty()); + assertFalse(supplier.getQueryEndPoint().isPresent()); + assertTrue(reconnect(newConnection())); + assertEquals(endpoints(0, 1, 2), supplier.getSeedNodes()); + advertised = endpoints(0, 1, 2); + refreshes.remove(0).run(); + assertEquals(advertised, supplier.get()); + } + + @Test + public void testAllUnavailableRetriesAreBoundedAndDeduplicated() throws Exception { + supplier.close(); + supplier = createSupplier(endpoints(0, 0, 1, 2)); + advertised = endpoints(1); + supplier.run(); + nodes.forEach(node -> node.online = false); + attempts.clear(); + assertFalse(reconnect(newConnection())); + assertEquals(SessionConfig.RETRY_NUM * 3, attempts.size()); + for (int i = 0; i < SessionConfig.RETRY_NUM; i++) { + assertEquals(endpoints(1, 0, 2), attempts.subList(i * 3, (i + 1) * 3)); + } + assertTrue(refreshes.isEmpty()); + } + + @Test + public void testSeedSnapshotIsIndependentOfCallerAndDiscovery() { + supplier.close(); + List seeds = endpoints(0, 1); + supplier = createSupplier(seeds); + seeds.get(0).setPort(1); + seeds.clear(); + advertised = endpoints(2); + supplier.run(); + assertEquals(endpoints(0, 1), supplier.getSeedNodes()); + assertEquals(endpoints(2), supplier.get()); + } + + @Test + public void testBackgroundRefreshDoesNotContactSeedsWhileCurrentNodeWorks() { + advertised = endpoints(1); + supplier.run(); + attempts.clear(); + supplier.run(); + assertTrue(attempts.isEmpty()); + assertEquals(endpoints(1), supplier.get()); + } + + @Test + public void testSessionRetainsInitialUrlsAfterDefaultEndpointChanges() throws Exception { + List urls = new ArrayList<>(Arrays.asList(url(0), url(1))); + Session session = new Session.Builder().nodeUrls(urls).build(); + urls.clear(); + session.defaultEndPoint = nodes.get(2).endpoint; + List seeds = Whitebox.invokeMethod(session, "getNodeUrls"); + assertEquals(2, seeds.size()); + assertTrue(seeds.containsAll(endpoints(0, 1))); + + Session singleNode = + new Session(nodes.get(0).endpoint.getIp(), nodes.get(0).endpoint.getPort()); + singleNode.defaultEndPoint = nodes.get(2).endpoint; + assertEquals(endpoints(0), Whitebox.invokeMethod(singleNode, "getNodeUrls")); + } + + @Test + public void testQuerySelectionDuringEmptyDiscoveryRefresh() throws Exception { + advertised = endpoints(0); + supplier.run(); + ExecutorService reader = Executors.newSingleThreadExecutor(); + AtomicBoolean finished = new AtomicBoolean(); + CountDownLatch started = new CountDownLatch(1); + Future selection = + reader.submit( + () -> { + started.countDown(); + while (!finished.get()) { + supplier + .getQueryEndPoint() + .ifPresent(endpoint -> assertEquals(nodes.get(0).endpoint, endpoint)); + } + }); + try { + assertTrue(started.await(5, TimeUnit.SECONDS)); + for (int i = 0; i < 100; i++) { + advertised = i % 2 == 0 ? Collections.emptyList() : endpoints(0); + supplier.run(); + } + } finally { + finished.set(true); + reader.shutdownNow(); + assertTrue(reader.awaitTermination(5, TimeUnit.SECONDS)); + } + selection.get(5, TimeUnit.SECONDS); + } + + @Test + public void testRefreshRequestsAreCoalescedAndIgnoredAfterClose() { + supplier.requestRefresh(); + supplier.requestRefresh(); + assertEquals(1, refreshes.size()); + refreshes.remove(0).run(); + supplier.requestRefresh(); + assertEquals(1, refreshes.size()); + supplier.close(); + refreshes.remove(0).run(); + supplier.requestRefresh(); + assertTrue(refreshes.isEmpty()); + } + + private NodesSupplier createSupplier(List seeds) { + return NodesSupplier.createNodeSupplier( + seeds, + executor, + "root", + "root", + ZoneId.of("UTC"), + SessionConfig.DEFAULT_INITIAL_BUFFER_CAPACITY, + SessionConfig.DEFAULT_MAX_FRAME_SIZE, + 1000, + false, + null, + null, + null, + null, + "TLS", + false, + "V_1_0"); + } + + private SessionConnection newConnection() { + SessionConnection connection = new SessionConnection("tree"); + Session session = new Session.Builder().nodeUrls(Collections.singletonList(url(0))).build(); + session.connectionTimeoutInMs = 1000; + Whitebox.setInternalState(connection, "session", session); + Whitebox.setInternalState(connection, "zoneId", ZoneId.of("UTC")); + Whitebox.setInternalState(connection, "transport", mock(TTransport.class)); + Whitebox.setInternalState(connection, "availableNodes", supplier); + connections.add(connection); + return connection; + } + + private boolean reconnect(SessionConnection connection) throws Exception { + return Whitebox.invokeMethod(connection, "reconnect"); + } + + private List endpoints(int... indices) { + List result = new ArrayList<>(); + for (int index : indices) { + result.add(nodes.get(index).endpoint.deepCopy()); + } + return result; + } + + private String url(int index) { + return UrlUtils.convertTEndPointIpv4AndIpv6Url(nodes.get(index).endpoint); + } + + private TSStatus success() { + return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()); + } + + private class TestNode implements AutoCloseable { + private final TEndPoint endpoint; + private final TServerSocket socket; + private final TThreadPoolServer server; + private final ExecutorService workers = Executors.newCachedThreadPool(); + private final Thread thread; + private volatile boolean online = true; + + private TestNode() throws Exception { + ServerSocket listener = new ServerSocket(0); + endpoint = new TEndPoint("127.0.0.1", listener.getLocalPort()); + socket = new TServerSocket(listener); + IClientRPCService.Iface handler = mock(IClientRPCService.Iface.class); + when(handler.openSession(any())) + .thenAnswer( + invocation -> { + attempts.add(endpoint); + checkOnline(); + return new TSOpenSessionResp(success(), TSProtocolVersion.IOTDB_SERVICE_PROTOCOL_V3) + .setSessionId(1); + }); + when(handler.requestStatementId(anyLong())).thenReturn(1L); + when(handler.closeSession(any())).thenReturn(success()); + when(handler.closeOperation(any())).thenReturn(success()); + when(handler.executeQueryStatementV2(any())) + .thenAnswer( + invocation -> { + checkOnline(); + TsBlockBuilder block = + new TsBlockBuilder(Arrays.asList(TSDataType.TEXT, TSDataType.INT32)); + for (TEndPoint address : advertised) { + block.getTimeColumnBuilder().writeLong(0); + block.getColumnBuilder(0).writeBinary(BytesUtils.valueOf(address.getIp())); + block.getColumnBuilder(1).writeInt(address.getPort()); + block.declarePosition(); + } + return new TSExecuteStatementResp(success()) + .setQueryId(1) + .setColumns(Arrays.asList("RpcAddress", "RpcPort")) + .setDataTypeList(Arrays.asList("TEXT", "INT32")) + .setQueryResult( + advertised.isEmpty() + ? Collections.emptyList() + : Collections.singletonList( + new TsBlockSerde().serialize(block.build()))) + .setIgnoreTimeStamp(true) + .setMoreData(false); + }); + server = + new TThreadPoolServer( + new TThreadPoolServer.Args(socket) + .processor(new IClientRPCService.Processor<>(handler)) + .protocolFactory(new TBinaryProtocol.Factory()) + .transportFactory(DeepCopyRpcTransportFactory.INSTANCE) + .executorService(workers)); + CountDownLatch ready = new CountDownLatch(1); + thread = + new Thread( + () -> { + ready.countDown(); + server.serve(); + }, + "seed-fallback-test-server"); + thread.setDaemon(true); + thread.start(); + assertTrue(ready.await(5, TimeUnit.SECONDS)); + } + + private void checkOnline() throws TException { + if (!online) { + throw new TException("Node is restarting"); + } + } + + @Override + public void close() throws Exception { + server.stop(); + socket.close(); + workers.shutdownNow(); + thread.join(5000); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowAvailableUrlsTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowAvailableUrlsTask.java index 6d957769b4428..9c1311a77fe0a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowAvailableUrlsTask.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowAvailableUrlsTask.java @@ -19,7 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.execution.config.metadata; -import org.apache.iotdb.commons.cluster.RegionStatus; +import org.apache.iotdb.commons.cluster.NodeStatus; import org.apache.iotdb.commons.schema.column.ColumnHeader; import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant; import org.apache.iotdb.confignode.rpc.thrift.TDataNodeInfo; @@ -58,7 +58,9 @@ public static void buildTsBlock( if (showDataNodesResp.getDataNodesInfoList() != null) { for (TDataNodeInfo dataNodeInfo : showDataNodesResp.getDataNodesInfoList()) { String status = dataNodeInfo.getStatus(); - if (RegionStatus.Removing.getStatus().equals(status)) { + if (NodeStatus.Removing.getStatus().equals(status) + || NodeStatus.Unknown.getStatus().equals(status) + || NodeStatus.Stopped.getStatus().equals(status)) { continue; } builder.getTimeColumnBuilder().writeLong(0L); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowAvailableUrlsTaskTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowAvailableUrlsTaskTest.java new file mode 100644 index 0000000000000..734915bcff9d4 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ShowAvailableUrlsTaskTest.java @@ -0,0 +1,78 @@ +/* + * 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.db.queryengine.plan.execution.config.metadata; + +import org.apache.iotdb.commons.cluster.NodeStatus; +import org.apache.iotdb.confignode.rpc.thrift.TDataNodeInfo; +import org.apache.iotdb.confignode.rpc.thrift.TShowDataNodesResp; +import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult; + +import com.google.common.util.concurrent.SettableFuture; +import org.apache.tsfile.read.common.block.TsBlock; +import org.junit.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +import static org.junit.Assert.assertEquals; + +public class ShowAvailableUrlsTaskTest { + @Test + public void testOnlyRunningAndReadOnlyNodesAreReturned() throws Exception { + List nodes = + Arrays.stream(NodeStatus.values()) + .map( + status -> + new TDataNodeInfo() + .setStatus(status.getStatus()) + .setStatusReason("reason") + .setRpcAddresss("127.0.0.1") + .setRpcPort(6667 + status.ordinal())) + .collect(Collectors.toList()); + TsBlock result = execute(nodes); + assertEquals(2, result.getPositionCount()); + assertEquals("127.0.0.1", result.getColumn(0).getBinary(0).toString()); + assertEquals(6667 + NodeStatus.Running.ordinal(), result.getColumn(1).getInt(0)); + assertEquals(6667 + NodeStatus.ReadOnly.ordinal(), result.getColumn(1).getInt(1)); + } + + @Test + public void testNoOnlineNodes() throws Exception { + assertEquals( + 0, + execute( + Arrays.asList( + new TDataNodeInfo().setStatus(NodeStatus.Unknown.getStatus()), + new TDataNodeInfo().setStatus(NodeStatus.Stopped.getStatus()), + new TDataNodeInfo().setStatus(NodeStatus.Removing.getStatus()))) + .getPositionCount()); + assertEquals(0, execute(Collections.emptyList()).getPositionCount()); + assertEquals(0, execute(null).getPositionCount()); + } + + private TsBlock execute(List nodes) throws Exception { + SettableFuture future = SettableFuture.create(); + ShowAvailableUrlsTask.buildTsBlock( + new TShowDataNodesResp().setDataNodesInfoList(nodes), future); + return future.get().getResultSet(); + } +}