From 6508e3a13424aafc683b8fea11175d10e567bf48 Mon Sep 17 00:00:00 2001 From: Yann Massard Date: Mon, 2 Nov 2026 23:17:13 +0000 Subject: [PATCH] Add insertTablets(List) and insertAlignedTablets(List) to Session and SessionPool The map-based insertTablets methods allow only one tablet per device, because the map is keyed by device id. The new list overloads accept several tablets of the same device and send them in one request. With redirection enabled, each tablet is routed by its own device id, so tablets of devices with the same leader go to that leader together. insertTabletsWithLeaderCache now takes a list of KeyedTablet instead of a map, so a leader cache key may occur several times. The map methods keep routing by their map keys. --- .../org/apache/iotdb/isession/ISession.java | 16 ++ .../iotdb/isession/pool/ISessionPool.java | 16 ++ .../org/apache/iotdb/session/Session.java | 95 +++++-- .../iotdb/session/pool/SessionPool.java | 57 +++++ .../session/SessionInsertTabletListTest.java | 233 ++++++++++++++++++ .../iotdb/session/pool/SessionPoolTest.java | 65 +++++ 6 files changed, 459 insertions(+), 23 deletions(-) create mode 100644 iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionInsertTabletListTest.java diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java index 70d4266d6c889..c909c6a78ad87 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java @@ -367,6 +367,22 @@ void insertAlignedTablets(Map tablets) void insertAlignedTablets(Map tablets, boolean sorted) throws IoTDBConnectionException, StatementExecutionException; + /** + * Inserts several tablets in one request. Unlike {@link #insertTablets(Map)}, several tablets may + * belong to the same device, e.g. one tablet per measurement for measurements with different + * timestamps. With redirection enabled, each tablet is routed by its own device id. + */ + void insertTablets(List tablets) + throws IoTDBConnectionException, StatementExecutionException; + + /** + * Inserts several tablets of aligned devices in one request. Unlike {@link + * #insertAlignedTablets(Map)}, several tablets may belong to the same device. With redirection + * enabled, each tablet is routed by its own device id. + */ + void insertAlignedTablets(List tablets) + throws IoTDBConnectionException, StatementExecutionException; + void testInsertTablet(Tablet tablet) throws IoTDBConnectionException, StatementExecutionException; void testInsertTablet(Tablet tablet, boolean sorted) diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/pool/ISessionPool.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/pool/ISessionPool.java index acf3a9bac913c..4b59c0c7fe4da 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/pool/ISessionPool.java +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/pool/ISessionPool.java @@ -73,6 +73,22 @@ void insertTablets(Map tablets, boolean sorted) void insertAlignedTablets(Map tablets, boolean sorted) throws IoTDBConnectionException, StatementExecutionException; + /** + * Inserts several tablets in one request. Unlike {@link #insertTablets(Map)}, several tablets may + * belong to the same device, e.g. one tablet per measurement for measurements with different + * timestamps. With redirection enabled, each tablet is routed by its own device id. + */ + void insertTablets(List tablets) + throws IoTDBConnectionException, StatementExecutionException; + + /** + * Inserts several tablets of aligned devices in one request. Unlike {@link + * #insertAlignedTablets(Map)}, several tablets may belong to the same device. With redirection + * enabled, each tablet is routed by its own device id. + */ + void insertAlignedTablets(List tablets) + throws IoTDBConnectionException, StatementExecutionException; + void insertRecords( List deviceIds, List times, 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 de3ee417afc0d..e4376fb44714f 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 @@ -3207,17 +3207,9 @@ public void insertTablets(Map tablets) public void insertTablets(Map tablets, boolean sorted) throws IoTDBConnectionException, StatementExecutionException { if (enableRedirection) { - insertTabletsWithLeaderCache(tablets, sorted, false); + insertTabletsWithLeaderCache(tabletMapToList(tablets), sorted, false); } else { - TSInsertTabletsReq request = - genTSInsertTabletsReq(new ArrayList<>(tablets.values()), sorted, false); - if (request == null) { - return; - } - try { - getDefaultSessionConnection().insertTablets(request); - } catch (RedirectException ignored) { - } + insertTabletsWithDefaultConnection(new ArrayList<>(tablets.values()), sorted, false); } } @@ -3246,29 +3238,86 @@ public void insertAlignedTablets(Map tablets) public void insertAlignedTablets(Map tablets, boolean sorted) throws IoTDBConnectionException, StatementExecutionException { if (enableRedirection) { - insertTabletsWithLeaderCache(tablets, sorted, true); + insertTabletsWithLeaderCache(tabletMapToList(tablets), sorted, true); } else { - TSInsertTabletsReq request = - genTSInsertTabletsReq(new ArrayList<>(tablets.values()), sorted, true); - if (request == null) { - return; - } - try { - getDefaultSessionConnection().insertTablets(request); - } catch (RedirectException ignored) { + insertTabletsWithDefaultConnection(new ArrayList<>(tablets.values()), sorted, true); + } + } + + /** + * Inserts several tablets of possibly the same device in one request. With redirection enabled, + * each tablet is routed by its own device id, so tablets whose devices have the same leader are + * sent to that leader together. + * + * @param tablets data batch, possibly several tablets per device + */ + @Override + public void insertTablets(List tablets) + throws IoTDBConnectionException, StatementExecutionException { + insertTabletList(tablets, false); + } + + /** + * Inserts several tablets of possibly the same aligned device in one request. With redirection + * enabled, each tablet is routed by its own device id, so tablets whose devices have the same + * leader are sent to that leader together. + * + * @param tablets data batch, possibly several tablets per device + */ + @Override + public void insertAlignedTablets(List tablets) + throws IoTDBConnectionException, StatementExecutionException { + insertTabletList(tablets, true); + } + + private void insertTabletList(List tablets, boolean isAligned) + throws IoTDBConnectionException, StatementExecutionException { + if (enableRedirection) { + List keyedTablets = new ArrayList<>(tablets.size()); + for (Tablet tablet : tablets) { + keyedTablets.add(new KeyedTablet(tablet.getDeviceId(), tablet)); } + insertTabletsWithLeaderCache(keyedTablets, false, isAligned); + } else { + insertTabletsWithDefaultConnection(tablets, false, isAligned); } } + private void insertTabletsWithDefaultConnection( + List tablets, boolean sorted, boolean isAligned) + throws IoTDBConnectionException, StatementExecutionException { + TSInsertTabletsReq request = genTSInsertTabletsReq(tablets, sorted, isAligned); + if (request == null) { + return; + } + try { + getDefaultSessionConnection().insertTablets(request); + } catch (RedirectException ignored) { + } + } + + /** A tablet and the key under which its leader is looked up in the leader cache. */ + private record KeyedTablet(String leaderCacheKey, Tablet tablet) {} + + /** Keys each tablet by its map key, which callers of the map methods may choose freely. */ + private static List tabletMapToList(Map tablets) { + return tablets.entrySet().stream() + .map(entry -> new KeyedTablet(entry.getKey(), entry.getValue())) + .toList(); + } + + /** + * @param tablets the tablets to insert; the same leader cache key may occur several times + */ private void insertTabletsWithLeaderCache( - Map tablets, boolean sorted, boolean isAligned) + List tablets, boolean sorted, boolean isAligned) throws IoTDBConnectionException, StatementExecutionException { Map tabletGroup = new HashMap<>(); - for (Entry entry : tablets.entrySet()) { - final SessionConnection connection = getSessionConnection(entry.getKey()); + for (KeyedTablet keyedTablet : tablets) { + final SessionConnection connection = getSessionConnection(keyedTablet.leaderCacheKey()); TSInsertTabletsReq request = tabletGroup.computeIfAbsent(connection, k -> new TSInsertTabletsReq()); - updateTSInsertTabletsReq(request, entry.getValue(), sorted, isAligned); + updateTSInsertTabletsReq(request, keyedTablet.tablet(), sorted, isAligned); } tabletGroup.entrySet().removeIf(e -> e.getValue().getPrefixPathsSize() == 0); 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 16b79c6415c4b..5903a44f46558 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 @@ -1113,6 +1113,63 @@ public void insertAlignedTablets(Map tablets, boolean sorted) } } + /** + * Inserts several tablets in one request; several tablets may belong to the same device. + * + * @param tablets data batch, possibly several tablets per device + * @see Session#insertTablets(List) + */ + @SuppressWarnings({"squid:S112"}) // ignore Generic exceptions should never be thrown + @Override + public void insertTablets(List tablets) + throws IoTDBConnectionException, StatementExecutionException { + ISession session = getSession(); + try { + session.insertTablets(tablets); + putBack(session); + } catch (IoTDBConnectionException e) { + // TException means the connection is broken, remove it and get a new one. + LOGGER.warn(SessionMessages.INSERT_TABLETS_FAILED, e); + cleanSessionAndMayThrowConnectionException(session, FINAL_RETRY, e); + } catch (StatementExecutionException | RuntimeException e) { + putBack(session); + throw e; + } catch (Throwable e) { + LOGGER.error(SessionMessages.UNEXPECTED_ERROR_IN_INSERT_TABLETS, e); + putBack(session); + throw new RuntimeException(e); + } + } + + /** + * Inserts several tablets of aligned devices in one request; several tablets may belong to the + * same device. + * + * @param tablets data batch, possibly several tablets per device + * @see Session#insertAlignedTablets(List) + */ + @SuppressWarnings({"squid:S112"}) // ignore Generic exceptions should never be thrown + @Override + public void insertAlignedTablets(List tablets) + throws IoTDBConnectionException, StatementExecutionException { + ISession session = getSession(); + try { + session.insertAlignedTablets(tablets); + putBack(session); + } catch (IoTDBConnectionException e) { + // TException means the connection is broken, remove it and get a new one. + LOGGER.warn(SessionMessages.INSERT_ALIGNED_TABLETS_FAILED, e); + cleanSessionAndMayThrowConnectionException(session, FINAL_RETRY, e); + } catch (StatementExecutionException | RuntimeException e) { + putBack(session); + throw e; + } catch (Throwable e) { + LOGGER.error(SessionMessages.UNEXPECTED_ERROR_IN_INSERT_ALIGNED_TABLETS, e); + putBack(session); + throw new RuntimeException(e); + } + } + /** * Insert data in batch format, which can reduce the overhead of network. This method is just like * jdbc batch insert, we pack some insert request in batch and send them to server If you want diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionInsertTabletListTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionInsertTabletListTest.java new file mode 100644 index 0000000000000..d002576a78bfc --- /dev/null +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionInsertTabletListTest.java @@ -0,0 +1,233 @@ +/* + * 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.isession.SessionConfig; +import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.rpc.RedirectException; +import org.apache.iotdb.rpc.StatementExecutionException; +import org.apache.iotdb.service.rpc.thrift.TSInsertTabletsReq; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.IMeasurementSchema; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.After; +import org.junit.Test; + +import java.time.ZoneId; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +/** {@link Session#insertTablets(List)} and {@link Session#insertAlignedTablets(List)}. */ +public class SessionInsertTabletListTest { + + private static final TEndPoint DEFAULT_ENDPOINT = new TEndPoint("127.0.0.1", 55560); + private static final TEndPoint LEADER_OF_D1 = new TEndPoint("127.0.0.1", 55561); + private static final TEndPoint LEADER_OF_D2 = new TEndPoint("127.0.0.1", 55562); + + private static final String D1 = "root.db.d1"; + private static final String D2 = "root.db.d2"; + + private static final Map LEADERS = new HashMap<>(); + + static { + LEADERS.put(D1, LEADER_OF_D1); + LEADERS.put(D2, LEADER_OF_D2); + } + + /** The requests each endpoint received, in order. */ + private final Map> requests = new LinkedHashMap<>(); + + private Session session; + + @After + public void tearDown() throws IoTDBConnectionException { + if (session != null) { + session.close(); + } + } + + @Test + public void test_InsertTablets_WhenLeadersAreUnknown_SendsAllTabletsToDefaultEndpoint() + throws IoTDBConnectionException, StatementExecutionException { + session = openSession(true); + + session.insertTablets(Arrays.asList(tablet(D1, "s1"), tablet(D1, "s2"), tablet(D2, "s1"))); + + assertEquals(Collections.singleton(DEFAULT_ENDPOINT), requests.keySet()); + assertEquals(1, requests.get(DEFAULT_ENDPOINT).size()); + assertEquals(Arrays.asList(D1, D1, D2), requests.get(DEFAULT_ENDPOINT).get(0).getPrefixPaths()); + // the redirect taught the session the leaders + assertEquals(LEADER_OF_D1, session.deviceIdToEndpoint.get(D1)); + assertEquals(LEADER_OF_D2, session.deviceIdToEndpoint.get(D2)); + } + + @Test + public void test_InsertTablets_WhenLeadersAreKnown_SendsOneRequestPerLeader() + throws IoTDBConnectionException, StatementExecutionException { + session = openSession(true); + session.insertTablets(Arrays.asList(tablet(D1, "s1"), tablet(D2, "s1"))); + requests.clear(); + + session.insertTablets( + Arrays.asList(tablet(D1, "s1"), tablet(D2, "s1"), tablet(D1, "s2"), tablet(D1, "s3"))); + + assertEquals(2, requests.size()); + assertEquals(1, requests.get(LEADER_OF_D1).size()); + TSInsertTabletsReq toLeaderOfD1 = requests.get(LEADER_OF_D1).get(0); + assertEquals(Arrays.asList(D1, D1, D1), toLeaderOfD1.getPrefixPaths()); + assertEquals( + Arrays.asList( + Collections.singletonList("s1"), + Collections.singletonList("s2"), + Collections.singletonList("s3")), + toLeaderOfD1.getMeasurementsList()); + assertEquals(1, requests.get(LEADER_OF_D2).size()); + assertEquals(Collections.singletonList(D2), requests.get(LEADER_OF_D2).get(0).getPrefixPaths()); + } + + @Test + public void test_InsertTablets_WhenRedirectionIsDisabled_SendsOneRequestToDefaultEndpoint() + throws IoTDBConnectionException, StatementExecutionException { + session = openSession(false); + + session.insertTablets(Arrays.asList(tablet(D1, "s1"), tablet(D1, "s2"), tablet(D2, "s1"))); + + assertEquals(Collections.singleton(DEFAULT_ENDPOINT), requests.keySet()); + assertEquals(1, requests.get(DEFAULT_ENDPOINT).size()); + TSInsertTabletsReq request = requests.get(DEFAULT_ENDPOINT).get(0); + assertEquals(Arrays.asList(D1, D1, D2), request.getPrefixPaths()); + assertFalse(request.isIsAligned()); + assertNull(session.deviceIdToEndpoint); + } + + @Test + public void test_InsertAlignedTablets_WhenLeadersAreKnown_SendsAlignedRequestPerLeader() + throws IoTDBConnectionException, StatementExecutionException { + session = openSession(true); + session.insertAlignedTablets(Arrays.asList(tablet(D1, "s1"), tablet(D2, "s1"))); + requests.clear(); + + session.insertAlignedTablets(Arrays.asList(tablet(D1, "s1"), tablet(D1, "s2"))); + + assertEquals(Collections.singleton(LEADER_OF_D1), requests.keySet()); + TSInsertTabletsReq request = requests.get(LEADER_OF_D1).get(0); + assertEquals(Arrays.asList(D1, D1), request.getPrefixPaths()); + assertTrue(request.isIsAligned()); + } + + @Test + public void test_InsertAlignedTablets_WhenRedirectionIsDisabled_SendsOneAlignedRequest() + throws IoTDBConnectionException, StatementExecutionException { + session = openSession(false); + + session.insertAlignedTablets(Arrays.asList(tablet(D1, "s1"), tablet(D1, "s2"))); + + assertEquals(Collections.singleton(DEFAULT_ENDPOINT), requests.keySet()); + TSInsertTabletsReq request = requests.get(DEFAULT_ENDPOINT).get(0); + assertEquals(Arrays.asList(D1, D1), request.getPrefixPaths()); + assertTrue(request.isIsAligned()); + } + + private Session openSession(boolean enableRedirection) throws IoTDBConnectionException { + Session newSession = new RecordingSession(enableRedirection); + newSession.open(); + return newSession; + } + + /** A tablet of one INT64 measurement with two rows. */ + private static Tablet tablet(String deviceId, String measurement) { + List schemas = new ArrayList<>(); + schemas.add(new MeasurementSchema(measurement, TSDataType.INT64)); + Tablet tablet = new Tablet(deviceId, schemas, 2); + for (int row = 0; row < 2; row++) { + tablet.addTimestamp(row, row); + tablet.addValue(measurement, row, (long) row); + } + return tablet; + } + + private class RecordingSession extends Session { + + RecordingSession(boolean enableRedirection) { + super( + DEFAULT_ENDPOINT.getIp(), + DEFAULT_ENDPOINT.getPort(), + SessionConfig.DEFAULT_USER, + SessionConfig.DEFAULT_PASSWORD, + SessionConfig.DEFAULT_FETCH_SIZE, + null, + SessionConfig.DEFAULT_INITIAL_BUFFER_CAPACITY, + SessionConfig.DEFAULT_MAX_FRAME_SIZE, + enableRedirection, + SessionConfig.DEFAULT_VERSION); + this.enableAutoFetch = false; + } + + @Override + public SessionConnection constructSessionConnection( + Session session, TEndPoint endpoint, ZoneId zoneId) { + return new RecordingSessionConnection(session, endpoint); + } + } + + /** + * Records every request and, like a DataNode, answers with a redirect if the request contains + * devices whose leader is another endpoint. + */ + private class RecordingSessionConnection extends SessionConnection { + + private final TEndPoint endPoint; + + RecordingSessionConnection(Session session, TEndPoint endPoint) { + super(session.sqlDialect); + this.endPoint = endPoint; + } + + @Override + public void close() {} + + @Override + protected void insertTablets(TSInsertTabletsReq request) throws RedirectException { + requests.computeIfAbsent(endPoint, k -> new ArrayList<>()).add(request); + Map redirects = new HashMap<>(); + for (String deviceId : request.getPrefixPaths()) { + if (!LEADERS.get(deviceId).equals(endPoint)) { + redirects.put(deviceId, LEADERS.get(deviceId)); + } + } + if (!redirects.isEmpty()) { + throw new RedirectException(redirects); + } + } + } +} diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/pool/SessionPoolTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/pool/SessionPoolTest.java index d08625cb978c4..d6c716ed20e3a 100644 --- a/iotdb-client/session/src/test/java/org/apache/iotdb/session/pool/SessionPoolTest.java +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/pool/SessionPoolTest.java @@ -293,6 +293,63 @@ public void testTestInsertTablet2() throws IoTDBConnectionException, StatementEx ((ConcurrentLinkedDeque) Whitebox.getInternalState(sessionPool, "queue")).size()); } + @Test + public void test_InsertTabletList_WhenSessionAcceptsIt_PassesTabletsOnAndPutsSessionBack() + throws IoTDBConnectionException, StatementExecutionException { + List tablets = Arrays.asList(createBooleanTablet(), createBooleanTablet()); + + sessionPool.insertTablets(tablets); + + Mockito.verify(session).insertTablets(tablets); + assertEquals(1, sessionPool.currentAvailableSize()); + assertEquals(0, sessionPool.currentOccupiedSize()); + } + + @Test + public void test_InsertAlignedTabletList_WhenSessionAcceptsIt_PassesTabletsOnAndPutsSessionBack() + throws IoTDBConnectionException, StatementExecutionException { + List tablets = Arrays.asList(createBooleanTablet(), createBooleanTablet()); + + sessionPool.insertAlignedTablets(tablets); + + Mockito.verify(session).insertAlignedTablets(tablets); + assertEquals(1, sessionPool.currentAvailableSize()); + assertEquals(0, sessionPool.currentOccupiedSize()); + } + + @Test + public void test_InsertTabletList_WhenSessionRejectsIt_RethrowsAndPutsSessionBack() + throws IoTDBConnectionException, StatementExecutionException { + List tablets = Collections.singletonList(createBooleanTablet()); + doThrow(new StatementExecutionException("rejected")).when(session).insertTablets(tablets); + + try { + sessionPool.insertTablets(tablets); + Assert.fail("expected StatementExecutionException"); + } catch (StatementExecutionException e) { + // expected + } + + assertEquals(1, sessionPool.currentAvailableSize()); + } + + @Test + public void test_InsertAlignedTabletList_WhenConnectionBreaks_ClosesSessionAndThrows() + throws IoTDBConnectionException, StatementExecutionException { + List tablets = Collections.singletonList(createBooleanTablet()); + doThrow(new IoTDBConnectionException("broken")).when(session).insertAlignedTablets(tablets); + + try { + sessionPool.insertAlignedTablets(tablets); + Assert.fail("expected IoTDBConnectionException"); + } catch (IoTDBConnectionException e) { + // expected + } + + Mockito.verify(session).close(); + assertEquals(0, sessionPool.currentAvailableSize()); + } + @Test public void testTestInsertTablets() throws IoTDBConnectionException, StatementExecutionException { List schemas = new ArrayList<>(); @@ -1634,4 +1691,12 @@ private List FakedFirstFetchTsBlockResult() { return Collections.singletonList(tsBlock); } + + private static Tablet createBooleanTablet() { + List schemas = new ArrayList<>(); + schemas.add(new MeasurementSchema("pressure", TSDataType.BOOLEAN)); + long[] timestamp = new long[] {1L, 2L}; + boolean[][] values = new boolean[][] {{true, false}}; + return new Tablet("device1", schemas, timestamp, values, new BitMap[1], 2); + } }