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); + } }