This is an automated email from the ASF dual-hosted git repository. Caideyipi pushed a commit to branch codex/cp-19064-dev13 in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 58453b30972c9e532e4e14f0e371bfb8c9a4b52a Author: Caideyipi <[email protected]> AuthorDate: Mon Nov 2 10:29:40 2026 +0800 Fix Pipe leader cache isolation across receivers (#19064) (cherry picked from commit 59ac348a3f2954cbdbc5e53c01a3cd475a187880) --- .../pipe/it/triple/AbstractPipeTripleManualIT.java | 6 +- .../iotdb/pipe/it/triple/IoTDBPipeForwardIT.java | 84 +++++++++++++++++ .../client/IoTDBDataNodeAsyncClientManager.java | 11 ++- .../IoTDBDataNodeCacheLeaderClientManager.java | 43 +++++++-- .../client/IoTDBDataNodeSyncClientManager.java | 12 ++- .../batch/PipeTransferBatchReqBuilder.java | 13 ++- .../thrift/async/IoTDBDataRegionAsyncSink.java | 3 +- .../thrift/sync/IoTDBDataRegionSyncSink.java | 3 +- .../sink/IoTDBDataNodeAsyncClientManagerTest.java | 103 +++++++++++++++++++++ .../batch/PipeTransferBatchReqBuilderTest.java | 53 +++++++++++ 10 files changed, 314 insertions(+), 17 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/AbstractPipeTripleManualIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/AbstractPipeTripleManualIT.java index c0cbec7f29c..a539f684fe1 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/AbstractPipeTripleManualIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/AbstractPipeTripleManualIT.java @@ -39,11 +39,15 @@ abstract class AbstractPipeTripleManualIT { env2 = MultiEnvFactory.getEnv(1); env3 = MultiEnvFactory.getEnv(2); setupConfig(); - env1.initClusterEnvironment(1, 1); + env1.initClusterEnvironment(1, getEnv1DataNodeCount()); env2.initClusterEnvironment(1, 1); env3.initClusterEnvironment(1, 1); } + protected int getEnv1DataNodeCount() { + return 1; + } + protected void setupConfig() { env1.getConfig() .getCommonConfig() diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBPipeForwardIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBPipeForwardIT.java index 25019a705f6..e169a21014c 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBPipeForwardIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBPipeForwardIT.java @@ -22,25 +22,109 @@ package org.apache.iotdb.pipe.it.triple; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.client.sync.SyncConfigNodeIServiceClient; import org.apache.iotdb.confignode.rpc.thrift.TCreatePipeReq; +import org.apache.iotdb.consensus.ConsensusFactory; import org.apache.iotdb.db.it.utils.TestUtils; +import org.apache.iotdb.isession.ISession; import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; import org.apache.iotdb.it.framework.IoTDBTestRunner; import org.apache.iotdb.itbase.category.MultiClusterIT3; +import org.apache.iotdb.itbase.env.BaseEnv; import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.MeasurementSchema; import org.junit.Assert; import org.junit.Test; import org.junit.experimental.categories.Category; import org.junit.runner.RunWith; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; +import java.util.List; import java.util.Map; +import java.util.stream.Collectors; @RunWith(IoTDBTestRunner.class) @Category({MultiClusterIT3.class}) public class IoTDBPipeForwardIT extends AbstractPipeTripleManualIT { + + @Override + protected int getEnv1DataNodeCount() { + return 3; + } + + @Override + protected void setupConfig() { + super.setupConfig(); + for (final BaseEnv env : Arrays.asList(env1, env2, env3)) { + env.getConfig() + .getCommonConfig() + .setPipeAutoSplitFullEnabled(true) + .setAutoCreateSchemaEnabled(true) + .setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS); + } + } + + @Test + public void testLocalWriteIsForwardedToDownstreamWithDefaultSourceParameters() throws Exception { + TestUtils.executeNonQueries( + env1, + Collections.singletonList( + "create pipe a2b with source('forwarding-pipe-requests'='false'," + + "'capture.tree'='true','capture.table'='true') with sink('node-urls'='" + + env2.getDataNodeWrapper(0).getIpAndPortString() + + "')"), + null); + TestUtils.executeNonQueries( + env2, + Arrays.asList( + "create pipe b2a with source('forwarding-pipe-requests'='false'," + + "'capture.tree'='true','capture.table'='true') with sink('node-urls'='" + + env1.getDataNodeWrapperList().stream() + .map(DataNodeWrapper::getIpAndPortString) + .collect(Collectors.joining(",")) + + "')", + "create pipe b2c with sink('node-urls'='" + + env3.getDataNodeWrapper(0).getIpAndPortString() + + "')"), + null); + + final List<MeasurementSchema> schemas = new ArrayList<>(); + for (int sensor = 0; sensor < 10; sensor++) { + schemas.add(new MeasurementSchema("s_" + sensor, TSDataType.INT32)); + } + try (final ISession session = env2.getSessionConnection()) { + for (int batch = 0; batch < 60; batch++) { + for (int device = 0; device < 2; device++) { + final Tablet tablet = new Tablet("root.test.sg_0.d_" + device, schemas, 1000); + for (int row = 0; row < 1000; row++) { + final long timestamp = 1577808000000L + (batch * 1000L + row) * 500; + tablet.addTimestamp(row, timestamp); + for (int sensor = 0; sensor < schemas.size(); sensor++) { + tablet.addValue("s_" + sensor, row, batch * 1000 + row); + } + } + tablet.rowSize = 1000; + session.insertAlignedTablet(tablet); + } + // Pace the writes so the workload continues after the source's initial flush. + Thread.sleep(500); + } + } + for (final BaseEnv env : Arrays.asList(env1, env2, env3)) { + TestUtils.assertDataEventuallyOnEnv( + env, + "select count(s_0) from root.test.sg_0.* align by device", + "Device,count(s_0),", + new HashSet<>(Arrays.asList("root.test.sg_0.d_0,60000,", "root.test.sg_0.d_1,60000,")), + 60); + } + } + @Test public void testForwardingPipeRequests() throws Exception { final DataNodeWrapper env2DataNode = env2.getDataNodeWrapper(0); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeAsyncClientManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeAsyncClientManager.java index 20d921e4672..5f1ff3be390 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeAsyncClientManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeAsyncClientManager.java @@ -73,6 +73,7 @@ public class IoTDBDataNodeAsyncClientManager extends IoTDBClientManager private static final Map<String, Integer> CLIENT_RESOURCE_REF_COUNT = new ConcurrentHashMap<>(); private final String receiverAttributes; private final String clientResourceKey; + private final String receiverKey; // clientResourceKey -> IClientManager<TEndPoint, AsyncPipeDataTransferServiceClient> private static final Map<String, IClientManager<TEndPoint, AsyncPipeDataTransferServiceClient>> @@ -114,6 +115,7 @@ public class IoTDBDataNodeAsyncClientManager extends IoTDBClientManager shouldMarkAsPipeRequest); endPointSet = new HashSet<>(endPoints); + receiverKey = IoTDBDataNodeCacheLeaderClientManager.generateReceiverKey(endPoints); receiverAttributes = String.format( @@ -180,7 +182,7 @@ public class IoTDBDataNodeAsyncClientManager extends IoTDBClientManager return borrowClient(); } - return borrowClient(LEADER_CACHE_MANAGER.getLeaderEndPoint(deviceId)); + return borrowClient(getLeaderEndPoint(deviceId)); } public AsyncPipeDataTransferServiceClient borrowClient(final TEndPoint endPoint) @@ -376,6 +378,11 @@ public class IoTDBDataNodeAsyncClientManager extends IoTDBClientManager } } + @Override + public TEndPoint getLeaderEndPoint(final String deviceId) { + return useLeaderCache ? LEADER_CACHE_MANAGER.getLeaderEndPoint(receiverKey, deviceId) : null; + } + public void updateLeaderCache(final String deviceId, final TEndPoint endPoint) { if (!useLeaderCache || deviceId == null || endPoint == null) { return; @@ -386,7 +393,7 @@ public class IoTDBDataNodeAsyncClientManager extends IoTDBClientManager endPointSet.add(endPoint); } - LEADER_CACHE_MANAGER.updateLeaderEndPoint(deviceId, endPoint); + LEADER_CACHE_MANAGER.updateLeaderEndPoint(receiverKey, deviceId, endPoint); } public ExecutorService getExecutor() { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeCacheLeaderClientManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeCacheLeaderClientManager.java index a9dbd43544e..b7bd9f71519 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeCacheLeaderClientManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeCacheLeaderClientManager.java @@ -27,22 +27,45 @@ import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryBlock; import com.github.benmanes.caffeine.cache.Cache; import com.github.benmanes.caffeine.cache.Caffeine; import com.github.benmanes.caffeine.cache.Weigher; +import org.apache.tsfile.utils.Pair; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.List; import java.util.concurrent.ConcurrentHashMap; +import java.util.stream.Collectors; public interface IoTDBDataNodeCacheLeaderClientManager { LeaderCacheManager LEADER_CACHE_MANAGER = new LeaderCacheManager(); + TEndPoint getLeaderEndPoint(String deviceId); + + static String generateReceiverKey(final List<TEndPoint> endPoints) { + return endPoints.stream() + .map( + endPoint -> { + final String host = endPoint.getIp(); + final String formattedHost = + host.contains(":") && !(host.startsWith("[") && host.endsWith("]")) + ? "[" + host + "]" + : host; + return formattedHost + ":" + endPoint.getPort(); + }) + .distinct() + .sorted() + .collect(Collectors.joining(",", "[", "]")); + } + class LeaderCacheManager { private static final Logger LOGGER = LoggerFactory.getLogger(LeaderCacheManager.class); private static final PipeConfig CONFIG = PipeConfig.getInstance(); // leader cache built by LRU - private final Cache<String, TEndPoint> device2endpoint; + // A device can have different leaders in different receiving clusters. Keep one bounded + // cache, but namespace its entries by the sink's configured endpoints, before redirection. + private final Cache<Pair<String, String>, TEndPoint> device2endpoint; // a hashmap to reuse the created endpoint private final ConcurrentHashMap<TEndPoint, TEndPoint> endPoints = new ConcurrentHashMap<>(); @@ -57,7 +80,10 @@ public interface IoTDBDataNodeCacheLeaderClientManager { device2endpoint = Caffeine.newBuilder() .maximumWeight(allocatedMemoryBlock.getMemoryUsageInBytes()) - .weigher((Weigher<String, TEndPoint>) (device, endPoint) -> device.getBytes().length) + .weigher( + (Weigher<Pair<String, String>, TEndPoint>) + (key, endPoint) -> + key.getLeft().getBytes().length + key.getRight().getBytes().length) .recordStats() .build(); @@ -95,20 +121,23 @@ public interface IoTDBDataNodeCacheLeaderClientManager { }); } - public TEndPoint getLeaderEndPoint(final String deviceId) { - return deviceId == null ? null : device2endpoint.getIfPresent(deviceId); + public TEndPoint getLeaderEndPoint(final String receiverKey, final String deviceId) { + return deviceId == null + ? null + : device2endpoint.getIfPresent(new Pair<>(receiverKey, deviceId)); } - public void updateLeaderEndPoint(final String deviceId, final TEndPoint endPoint) { + public void updateLeaderEndPoint( + final String receiverKey, final String deviceId, final TEndPoint endPoint) { if (deviceId == null || endPoint == null) { return; } final TEndPoint endPointFromMap = endPoints.putIfAbsent(endPoint, endPoint); if (endPointFromMap != null) { - device2endpoint.put(deviceId, endPointFromMap); + device2endpoint.put(new Pair<>(receiverKey, deviceId), endPointFromMap); } else { - device2endpoint.put(deviceId, endPoint); + device2endpoint.put(new Pair<>(receiverKey, deviceId), endPoint); } } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeSyncClientManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeSyncClientManager.java index 2a05fe61378..d2394120f43 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeSyncClientManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeSyncClientManager.java @@ -42,6 +42,8 @@ public class IoTDBDataNodeSyncClientManager extends IoTDBSyncClientManager private static final Logger LOGGER = LoggerFactory.getLogger(IoTDBDataNodeSyncClientManager.class); + private final String receiverKey; + public IoTDBDataNodeSyncClientManager( final List<TEndPoint> endPoints, final String username, @@ -68,6 +70,7 @@ public class IoTDBDataNodeSyncClientManager extends IoTDBSyncClientManager loadTsFileStrategy, validateTsFile, shouldMarkAsPipeRequest); + receiverKey = IoTDBDataNodeCacheLeaderClientManager.generateReceiverKey(endPoints); } @Override @@ -88,7 +91,7 @@ public class IoTDBDataNodeSyncClientManager extends IoTDBSyncClientManager } public Pair<IoTDBSyncClient, Boolean> getClient(final String deviceId) { - final TEndPoint endPoint = LEADER_CACHE_MANAGER.getLeaderEndPoint(deviceId); + final TEndPoint endPoint = getLeaderEndPoint(deviceId); return useLeaderCache && endPoint != null && endPoint2ClientAndStatus.containsKey(endPoint) @@ -108,6 +111,11 @@ public class IoTDBDataNodeSyncClientManager extends IoTDBSyncClientManager : getClient(); } + @Override + public TEndPoint getLeaderEndPoint(final String deviceId) { + return useLeaderCache ? LEADER_CACHE_MANAGER.getLeaderEndPoint(receiverKey, deviceId) : null; + } + public void updateLeaderCache(final String deviceId, final TEndPoint endPoint) { if (!useLeaderCache || deviceId == null || endPoint == null) { return; @@ -120,7 +128,7 @@ public class IoTDBDataNodeSyncClientManager extends IoTDBSyncClientManager reconstructClient(endPoint); } - LEADER_CACHE_MANAGER.updateLeaderEndPoint(deviceId, endPoint); + LEADER_CACHE_MANAGER.updateLeaderEndPoint(receiverKey, deviceId, endPoint); } catch (final Exception e) { LOGGER.warn( "Failed to update leader cache for device {} with endpoint {}:{}.", diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilder.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilder.java index 3921d42e422..bd9627fdd24 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilder.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilder.java @@ -24,7 +24,6 @@ import org.apache.iotdb.commons.pipe.agent.task.progress.CommitterKey; import org.apache.iotdb.commons.pipe.event.EnrichedEvent; import org.apache.iotdb.db.pipe.event.common.tablet.PipeInsertNodeTabletInsertionEvent; import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent; -import org.apache.iotdb.db.pipe.sink.client.IoTDBDataNodeCacheLeaderClientManager; import org.apache.iotdb.db.storageengine.dataregion.wal.exception.WALPipeException; import org.apache.iotdb.metrics.impl.DoNothingHistogram; import org.apache.iotdb.metrics.type.Histogram; @@ -43,6 +42,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.function.Function; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_FORMAT_HYBRID_VALUE; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_FORMAT_KEY; @@ -68,6 +68,7 @@ public class PipeTransferBatchReqBuilder implements AutoCloseable { private final boolean usingTsFileBatch; private final boolean useLeaderCache; + private final Function<String, TEndPoint> leaderEndPointResolver; private final int requestMaxDelayInMs; private final long requestMaxBatchSizeInBytes; @@ -88,6 +89,13 @@ public class PipeTransferBatchReqBuilder implements AutoCloseable { private final Map<TEndPoint, PipeTabletEventPlainBatch> endPointToBatch = new HashMap<>(); public PipeTransferBatchReqBuilder(final PipeParameters parameters) { + // Transports without leader redirection, such as air gap, always use the default batch. + this(parameters, deviceId -> null); + } + + public PipeTransferBatchReqBuilder( + final PipeParameters parameters, final Function<String, TEndPoint> leaderEndPointResolver) { + this.leaderEndPointResolver = leaderEndPointResolver; usingTsFileBatch = parameters .getStringOrDefault( @@ -167,8 +175,7 @@ public class PipeTransferBatchReqBuilder implements AutoCloseable { return; } - final TEndPoint endPoint = - IoTDBDataNodeCacheLeaderClientManager.LEADER_CACHE_MANAGER.getLeaderEndPoint(deviceId); + final TEndPoint endPoint = leaderEndPointResolver.apply(deviceId); if (Objects.isNull(endPoint)) { defaultBatch.onEvent(event); return; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java index e6221f6706b..bb405c3966e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java @@ -208,7 +208,8 @@ public class IoTDBDataRegionAsyncSink extends IoTDBSink { isSplitTSFileBatchModeEnabled); if (isTabletBatchModeEnabled) { - tabletBatchBuilder = new PipeTransferBatchReqBuilder(parameters); + tabletBatchBuilder = + new PipeTransferBatchReqBuilder(parameters, clientManager::getLeaderEndPoint); } enableSendTsFileLimit = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java index 549751a0db1..367ba4e9666 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java @@ -100,7 +100,8 @@ public class IoTDBDataRegionSyncSink extends IoTDBDataNodeSyncSink { // tablet batch mode configuration if (isTabletBatchModeEnabled) { - tabletBatchBuilder = new PipeTransferBatchReqBuilder(parameters); + tabletBatchBuilder = + new PipeTransferBatchReqBuilder(parameters, clientManager::getLeaderEndPoint); } enableSendTsFileLimit = diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/IoTDBDataNodeAsyncClientManagerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/IoTDBDataNodeAsyncClientManagerTest.java index 423361002cd..5f66f50aaaa 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/IoTDBDataNodeAsyncClientManagerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/IoTDBDataNodeAsyncClientManagerTest.java @@ -20,17 +20,90 @@ package org.apache.iotdb.db.pipe.sink; import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.client.IClientManager; +import org.apache.iotdb.commons.client.async.AsyncPipeDataTransferServiceClient; import org.apache.iotdb.db.pipe.sink.client.IoTDBDataNodeAsyncClientManager; +import org.apache.iotdb.db.pipe.sink.client.IoTDBDataNodeSyncClientManager; import org.junit.Assert; import org.junit.Test; +import org.mockito.Mockito; import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; +import java.util.List; import java.util.concurrent.ExecutorService; public class IoTDBDataNodeAsyncClientManagerTest { + @Test + public void testAsyncBorrowDoesNotUseLeaderFromAnotherReceiver() throws Exception { + final TEndPoint firstReceiver = new TEndPoint("127.0.0.1", 10667); + final TEndPoint secondReceiver = new TEndPoint("127.0.0.1", 11667); + final TEndPoint firstLeader = new TEndPoint("127.0.0.1", 10668); + final IoTDBDataNodeAsyncClientManager firstManager = + createAsyncManager(new ArrayList<>(Collections.singletonList(firstReceiver))); + final IoTDBDataNodeAsyncClientManager secondManager = + createAsyncManager(new ArrayList<>(Collections.singletonList(secondReceiver))); + + final IClientManager<TEndPoint, AsyncPipeDataTransferServiceClient> clients = + Mockito.mock(IClientManager.class); + final AsyncPipeDataTransferServiceClient client = + Mockito.mock(AsyncPipeDataTransferServiceClient.class); + Mockito.when(clients.borrowClient(Mockito.any())).thenReturn(client); + Mockito.when(client.isHandshakeFinished()).thenReturn(true); + final Field clientManagerField = + IoTDBDataNodeAsyncClientManager.class.getDeclaredField("endPoint2Client"); + clientManagerField.setAccessible(true); + clientManagerField.set(secondManager, clients); + + try { + firstManager.updateLeaderCache("root.test.shared", firstLeader); + Assert.assertSame(client, secondManager.borrowClient("root.test.shared")); + Mockito.verify(clients).borrowClient(secondReceiver); + Mockito.verify(clients, Mockito.never()).borrowClient(firstLeader); + } finally { + firstManager.close(); + secondManager.close(); + } + } + + @Test + public void testLeaderCacheIsScopedToConfiguredReceiverEndPoints() { + final TEndPoint firstReceiver = new TEndPoint("127.0.0.1", 12667); + final TEndPoint secondReceiver = new TEndPoint("127.0.0.1", 12668); + final TEndPoint otherReceiver = new TEndPoint("::1", 13667); + final TEndPoint firstLeader = new TEndPoint("127.0.0.1", 12669); + final TEndPoint otherLeader = new TEndPoint("::1", 13668); + final String device = "root.test.scoped"; + final IoTDBDataNodeAsyncClientManager asyncManager = + createAsyncManager(new ArrayList<>(Arrays.asList(firstReceiver, secondReceiver))); + final TestIoTDBDataNodeSyncClientManager sameReceiverManager = + new TestIoTDBDataNodeSyncClientManager( + new ArrayList<>(Arrays.asList(secondReceiver, firstReceiver, firstReceiver))); + final TestIoTDBDataNodeSyncClientManager otherManager = + new TestIoTDBDataNodeSyncClientManager( + new ArrayList<>(Collections.singletonList(otherReceiver))); + + try { + asyncManager.updateLeaderCache(device, firstLeader); + Assert.assertEquals(firstLeader, asyncManager.getLeaderEndPoint(device)); + Assert.assertEquals(firstLeader, sameReceiverManager.getLeaderEndPoint(device)); + Assert.assertNull(otherManager.getLeaderEndPoint(device)); + + otherManager.updateLeaderCache(device, otherLeader); + Assert.assertEquals(otherLeader, otherManager.getLeaderEndPoint(device)); + Assert.assertEquals(firstLeader, asyncManager.getLeaderEndPoint(device)); + Assert.assertEquals(firstLeader, sameReceiverManager.getLeaderEndPoint(device)); + } finally { + asyncManager.close(); + sameReceiverManager.close(); + otherManager.close(); + } + } + @Test public void testClientResourcesShouldDifferentiateEndPoints() throws Exception { final IoTDBDataNodeAsyncClientManager firstManager = @@ -71,6 +144,12 @@ public class IoTDBDataNodeAsyncClientManagerTest { } } + private static IoTDBDataNodeAsyncClientManager createAsyncManager( + final List<TEndPoint> endPoints) { + return new IoTDBDataNodeAsyncClientManager( + endPoints, true, "round-robin", "user", "password", true, "sync", true, true, false); + } + private static String getReceiverAttributes(final IoTDBDataNodeAsyncClientManager manager) throws Exception { final Field field = @@ -99,4 +178,28 @@ public class IoTDBDataNodeAsyncClientManagerTest { field.setAccessible(true); return (ExecutorService) field.get(manager); } + + private static class TestIoTDBDataNodeSyncClientManager extends IoTDBDataNodeSyncClientManager { + + private TestIoTDBDataNodeSyncClientManager(final List<TEndPoint> endPoints) { + super( + endPoints, + "user", + "password", + false, + null, + null, + true, + "round-robin", + true, + "sync", + true, + true); + } + + @Override + protected void reconstructClient(final TEndPoint endPoint) { + // No client connection is needed to verify receiver-key isolation. + } + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilderTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilderTest.java index bfa8f9cd9aa..45154ca0010 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilderTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilderTest.java @@ -24,6 +24,7 @@ import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalExc import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent; import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager; import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryManager; +import org.apache.iotdb.db.pipe.sink.client.IoTDBDataNodeAsyncClientManager; import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters; import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent; @@ -34,6 +35,7 @@ import org.apache.tsfile.write.schema.MeasurementSchema; import org.junit.Assert; import org.junit.Test; +import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -41,6 +43,57 @@ import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CON public class PipeTransferBatchReqBuilderTest { + @Test + public void testBatchesDoNotUseLeaderFromAnotherReceiver() throws Exception { + final TEndPoint firstReceiver = new TEndPoint("127.0.0.1", 14667); + final TEndPoint secondReceiver = new TEndPoint("127.0.0.1", 15667); + final TEndPoint firstLeader = new TEndPoint("127.0.0.1", 14668); + final TEndPoint secondLeader = new TEndPoint("127.0.0.1", 15668); + final IoTDBDataNodeAsyncClientManager firstManager = createClientManager(firstReceiver); + final IoTDBDataNodeAsyncClientManager secondManager = createClientManager(secondReceiver); + final PipeParameters parameters = + new PipeParameters(Collections.singletonMap(CONNECTOR_IOTDB_BATCH_DELAY_MS_KEY, "0")); + final String deviceId = "root.test.device"; + + try (final PipeTransferBatchReqBuilder firstBuilder = + new PipeTransferBatchReqBuilder(parameters, firstManager::getLeaderEndPoint); + final PipeTransferBatchReqBuilder secondBuilder = + new PipeTransferBatchReqBuilder(parameters, secondManager::getLeaderEndPoint)) { + firstManager.updateLeaderCache(deviceId, firstLeader); + firstBuilder.onEvent(createEvent(1)); + secondBuilder.onEvent(createEvent(2)); + Assert.assertEquals( + firstLeader, firstBuilder.getAllNonEmptyAndShouldEmitBatches().get(0).getLeft()); + Assert.assertNull(secondBuilder.getAllNonEmptyAndShouldEmitBatches().get(0).getLeft()); + + secondManager.updateLeaderCache(deviceId, secondLeader); + secondBuilder.onEvent(createEvent(3)); + final List<Pair<TEndPoint, PipeTabletEventBatch>> secondBatches = + secondBuilder.getAllNonEmptyAndShouldEmitBatches(); + Assert.assertEquals(2, secondBatches.size()); + Assert.assertNull(secondBatches.get(0).getLeft()); + Assert.assertEquals(secondLeader, secondBatches.get(1).getLeft()); + Assert.assertEquals(firstLeader, firstManager.getLeaderEndPoint(deviceId)); + } finally { + firstManager.close(); + secondManager.close(); + } + } + + private static IoTDBDataNodeAsyncClientManager createClientManager(final TEndPoint receiver) { + return new IoTDBDataNodeAsyncClientManager( + new ArrayList<>(Collections.singletonList(receiver)), + true, + "round-robin", + "user", + "password", + true, + "sync", + true, + true, + false); + } + @Test public void testDetachedBatchCannotClearSubsequentEvent() throws Exception { final PipeTransferBatchReqBuilder builder =
