This is an automated email from the ASF dual-hosted git repository. jt2594838 pushed a commit to branch insert_numeric_row in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 67af28c728a79396cf2871d8f2b689d36ab6fcdb Author: Tian Jiang <[email protected]> AuthorDate: Thu Sep 17 16:04:31 2026 +0800 Add redirection --- .../it/session/IoTDBTableModelSessionIT.java | 70 ++++++- .../org/apache/iotdb/isession/ITableSession.java | 6 + .../isession/NumericRowPreparedStatement.java | 47 +++-- .../java/org/apache/iotdb/session/Session.java | 5 +- .../apache/iotdb/session/SessionConnection.java | 12 +- .../org/apache/iotdb/session/TableSession.java | 95 +++++++-- .../session/PreparedNumericRowRedirectionTest.java | 219 +++++++++++++++++++++ 7 files changed, 414 insertions(+), 40 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBTableModelSessionIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBTableModelSessionIT.java index 59169cb8554..db971c7cc02 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBTableModelSessionIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/IoTDBTableModelSessionIT.java @@ -43,11 +43,13 @@ import org.junit.Test; import org.junit.experimental.categories.Category; import org.junit.runner.RunWith; +import java.lang.reflect.Field; import java.time.LocalDate; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -319,10 +321,11 @@ public class IoTDBTableModelSessionIT { assertTrue(schemaError.getMessage().contains("v")); session.executeNonQueryStatement("DROP TABLE readings"); session.executeNonQueryStatement("CREATE TABLE readings(tag STRING TAG, v DOUBLE FIELD)"); - StatementExecutionException closedError = + IoTDBConnectionException closedError = assertThrows( - StatementExecutionException.class, + IoTDBConnectionException.class, () -> session.executePreparedInsertNumericRow(prepared, 2, values, tags)); + assertTrue(prepared.isClosed()); assertTrue(closedError.getMessage().contains("closed")); session.executePreparedInsertNumericRow(other, 3, values, tags); try (NumericRowPreparedStatement replacement = @@ -363,6 +366,69 @@ public class IoTDBTableModelSessionIT { } } + // Pin clients to every DataNode so at least one must learn a remote leader in a cluster. + @Test + public void testPreparedNumericRowRedirection() throws Exception { + try (ITableSession session = EnvFactory.getEnv().getTableSessionConnection()) { + session.executeNonQueryStatement("CREATE DATABASE numeric_redirect"); + session.executeNonQueryStatement("USE numeric_redirect"); + session.executeNonQueryStatement( + "CREATE TABLE readings(tag STRING TAG, v INT64 FIELD, n INT64 FIELD)"); + int clientIndex = 0; + boolean redirected = false; + for (var node : EnvFactory.getEnv().getDataNodeWrapperList()) { + try (ITableSession pinned = + EnvFactory.getEnv().getTableSessionConnection(List.of(node.getIpAndPortString()))) { + pinned.executeNonQueryStatement("USE numeric_redirect"); + redirected |= writeRedirectedNumericRows(pinned, clientIndex++ * 1000L); + } + } + try (ITableSessionPool pool = EnvFactory.getEnv().getTableSessionPool(1, "numeric_redirect"); + ITableSession pooled = pool.getSession()) { + redirected |= writeRedirectedNumericRows(pooled, clientIndex++ * 1000L); + } + if (EnvFactory.getEnv().getDataNodeWrapperList().size() > 1) { + assertTrue("Expected a prepared handle on a remote leader", redirected); + } + session.executeNonQueryStatement("FLUSH"); + try (SessionDataSet result = + session.executeQueryStatement("SELECT count(v),count(n) FROM readings")) { + assertTrue(result.hasNext()); + RowRecord counts = result.next(); + assertEquals(clientIndex * 32L, counts.getFields().get(0).getLongV()); + assertEquals(clientIndex * 16L, counts.getFields().get(1).getLongV()); + assertFalse(result.hasNext()); + } + } + } + + private boolean writeRedirectedNumericRows(ITableSession session, long start) throws Exception { + try (NumericRowPreparedStatement prepared = + session.prepareInsertNumericRow( + "readings", + List.of("tag", "v", "n"), + List.of(TSDataType.STRING, TSDataType.INT64, TSDataType.INT64), + List.of(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD))) { + BitMap nulls = new BitMap(2); + nulls.mark(1); + for (int round = 0; round < 4; round++) { + for (int device = 0; device < 8; device++) { + long time = start + round * 8 + device; + session.executePreparedInsertNumericRow( + prepared, + time, + new long[] {time, time}, + new String[] {"d" + device}, + round % 2 == 0 ? nulls : null); + } + } + // Count actual physical bindings: row counts alone cannot prove the client followed a hint. + Field bindings = NumericRowPreparedStatement.class.getDeclaredField("connectionBindings"); + bindings.setAccessible(true); + return ((Map<?, ?>) bindings.get(prepared)).size() > 1; + } + } + @Test public void testPreparedNumericRowAndTabletWritePerformance() throws Exception { Assume.assumeTrue( diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java index 35a2f71b7dc..1cfb0b86eed 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ITableSession.java @@ -70,6 +70,12 @@ public interface ITableSession extends AutoCloseable { List<ColumnCategory> columnCategories) throws StatementExecutionException, IoTDBConnectionException; + /** + * Executes a prepared row, following write redirection recommendations for subsequent rows. + * Destination handles are prepared lazily and released when the statement is closed. An execution + * error reported by the server closes the statement; prepare it again before reuse. Transport + * failures are propagated without replaying a write whose outcome is unknown. + */ void executePreparedInsertNumericRow( NumericRowPreparedStatement prepared, long timestamp, long[] values, String[] tagValues) throws StatementExecutionException, IoTDBConnectionException; diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/NumericRowPreparedStatement.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/NumericRowPreparedStatement.java index ee0b0f21804..2fa9ee21257 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/NumericRowPreparedStatement.java +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/NumericRowPreparedStatement.java @@ -12,7 +12,9 @@ import org.apache.tsfile.enums.ColumnCategory; import org.apache.tsfile.enums.TSDataType; import java.nio.ByteBuffer; +import java.util.IdentityHashMap; import java.util.List; +import java.util.Map; /** * Immutable metadata and reusable serialization buffers for repeated numeric table row inserts. @@ -22,6 +24,7 @@ import java.util.List; */ public final class NumericRowPreparedStatement implements AutoCloseable { private final String tableName; + private final String database; private final List<String> columnNames; private final List<TSDataType> dataTypes; private final List<ColumnCategory> columnCategories; @@ -30,17 +33,18 @@ public final class NumericRowPreparedStatement implements AutoCloseable { private ByteBuffer tagBuffer; private ByteBuffer valueBuffer; private ByteBuffer nullBitmapBuffer; - private long preparedId = -1; private Object owner; - private Runnable closeAction = () -> {}; + private final Map<Object, ConnectionBinding> connectionBindings = new IdentityHashMap<>(); private boolean closed; public NumericRowPreparedStatement( String tableName, List<String> columnNames, List<TSDataType> dataTypes, - List<ColumnCategory> columnCategories) { + List<ColumnCategory> columnCategories, + String database) { this.tableName = tableName; + this.database = database; this.columnNames = List.copyOf(columnNames); this.dataTypes = List.copyOf(dataTypes); this.columnCategories = List.copyOf(columnCategories); @@ -69,6 +73,10 @@ public final class NumericRowPreparedStatement implements AutoCloseable { return tableName; } + public String getDatabase() { + return database; + } + public List<String> getColumnNames() { return columnNames; } @@ -105,10 +113,6 @@ public final class NumericRowPreparedStatement implements AutoCloseable { return valueBuffer; } - public long getPreparedId() { - return preparedId; - } - /** Returns a reusable FIELD null bitmap buffer, allocated only for nullable executions. */ public ByteBuffer prepareNullBitmapBuffer() { if (nullBitmapBuffer == null) { @@ -122,21 +126,40 @@ public final class NumericRowPreparedStatement implements AutoCloseable { return owner; } - public void bind(long id, Object owner, Runnable closeAction) { - this.preparedId = id; + public void bind(long id, Object owner, long sessionId, Runnable closeAction) { this.owner = owner; - this.closeAction = closeAction; + bindConnection(owner, sessionId, id, closeAction); this.closed = false; } + /** Handles belong to a physical session, including its generation after a reconnect. */ + public Long getPreparedId(Object connection, long sessionId) { + ConnectionBinding binding = connectionBindings.get(connection); + return binding != null && binding.sessionId == sessionId ? binding.preparedId : null; + } + + public void bindConnection(Object connection, long sessionId, long id, Runnable closeAction) { + connectionBindings.put(connection, new ConnectionBinding(sessionId, id, closeAction)); + } + + private record ConnectionBinding(long sessionId, long preparedId, Runnable closeAction) {} + public boolean isClosed() { return closed; } @Override public void close() { - closeAction.run(); - closeAction = () -> {}; + if (closed) { + return; + } closed = true; + try { + for (ConnectionBinding binding : connectionBindings.values()) { + binding.closeAction.run(); + } + } finally { + connectionBindings.clear(); + } } } 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 f5c6ebd1b80..11a545da714 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 @@ -1357,8 +1357,7 @@ public class Session implements ISession { } } - private SessionConnection getSessionConnection(IDeviceID deviceId) - throws IoTDBConnectionException { + SessionConnection getSessionConnection(IDeviceID deviceId) throws IoTDBConnectionException { TEndPoint endPoint; if (enableRedirection && tableModelDeviceIdToEndpoint != null @@ -1440,7 +1439,7 @@ public class Session implements ISession { } } - private void handleRedirection(IDeviceID deviceId, TEndPoint endpoint) { + void handleRedirection(IDeviceID deviceId, TEndPoint endpoint) { if (enableRedirection) { // no need to redirection if (UrlUtils.isWildcardAddress(endpoint.ip)) { 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 f98e58965ae..28a75d0c307 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 @@ -781,20 +781,22 @@ public class SessionConnection { protected void executePreparedNumericRow( long id, long timestamp, ByteBuffer values, ByteBuffer tags) - throws IoTDBConnectionException, StatementExecutionException { + throws IoTDBConnectionException, StatementExecutionException, RedirectException { executePreparedNumericRow(id, timestamp, values, tags, null); } protected void executePreparedNumericRow( long id, long timestamp, ByteBuffer values, ByteBuffer tags, ByteBuffer nullBitmap) - throws IoTDBConnectionException, StatementExecutionException { + throws IoTDBConnectionException, StatementExecutionException, RedirectException { TSExecutePreparedNumericRowReq request = new TSExecutePreparedNumericRowReq(sessionId, id, timestamp, values, tags); if (nullBitmap != null) { request.setNullBitmap(nullBitmap); } try { - RpcUtils.verifySuccess(client.executePreparedNumericRow(request)); + // A redirect is a recommendation for later writes: this row has already been accepted. + // Never reconnect/replay a session-local prepared ID after an ambiguous transport failure. + RpcUtils.verifySuccessWithRedirection(client.executePreparedNumericRow(request)); } catch (TException e) { throw new IoTDBConnectionException(e); } @@ -810,6 +812,10 @@ public class SessionConnection { } } + long getSessionId() { + return sessionId; + } + private TSStatus insertRecordInternal(TSInsertRecordReq request) throws TException { request.setSessionId(sessionId); return client.insertRecord(request); diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSession.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSession.java index da93c37586f..9d3e6d8b9bb 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSession.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSession.java @@ -24,6 +24,7 @@ import org.apache.iotdb.isession.NumericRow; import org.apache.iotdb.isession.NumericRowPreparedStatement; import org.apache.iotdb.isession.SessionDataSet; 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.TSPrepareNumericRowReq; import org.apache.iotdb.session.i18n.SessionMessages; @@ -31,6 +32,7 @@ import org.apache.iotdb.session.util.SessionUtils; import org.apache.tsfile.enums.ColumnCategory; import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.utils.BitMap; import org.apache.tsfile.utils.ReadWriteIOUtils; import org.apache.tsfile.write.record.Tablet; @@ -38,6 +40,7 @@ import org.apache.tsfile.write.record.Tablet; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.List; +import java.util.Objects; public class TableSession implements ITableSession { @@ -84,18 +87,14 @@ public class TableSession implements ITableSession { .setDataTypes(types.stream().map(TSDataType::ordinal).toList()) .setColumnCategories(categories.stream().map(c -> (byte) c.ordinal()).toList()); NumericRowPreparedStatement prepared = - new NumericRowPreparedStatement(tableName, names, types, categories); + new NumericRowPreparedStatement(tableName, names, types, categories, session.getDatabase()); SessionConnection connection = session.getDefaultSessionConnection(); + long id = connection.prepareNumericRow(request); prepared.bind( - connection.prepareNumericRow(request), + id, connection, - () -> { - try { - connection.closePreparedNumericRow(prepared.getPreparedId()); - } catch (IoTDBConnectionException ignored) { - // Closing the session releases all server-side prepared rows. - } - }); + connection.getSessionId(), + closePreparedRow(connection, connection.getSessionId(), id)); return prepared; } @@ -115,7 +114,10 @@ public class TableSession implements ITableSession { BitMap nullBitmap) throws StatementExecutionException, IoTDBConnectionException { SessionConnection connection = session.getDefaultSessionConnection(); - if (prepared == null || prepared.isClosed() || prepared.getOwner() != connection) { + if (prepared == null + || prepared.isClosed() + || prepared.getOwner() != connection + || prepared.getPreparedId(connection, connection.getSessionId()) == null) { throw new IoTDBConnectionException( SessionMessages .EXCEPTION_PREPARED_NUMERIC_ROW_IS_CLOSED_OR_BELONGS_TO_ANOTHER_CONNECTION_PREPARE_IT_AGAIN_055CC7BC); @@ -130,6 +132,34 @@ public class TableSession implements ITableSession { SessionMessages.EXCEPTION_PREPARED_NUMERIC_ROW_VALUE_COUNT_MISMATCH_9ACD3A0C); } SessionUtils.validateNumericRowNullBitmap(nullBitmap, fieldCount); + // A prepared handle remains bound to its original database after USE. Do not prepare it on + // a redirected session whose database has changed in the meantime. + IDeviceID deviceId = null; + if (session.isEnableRedirection() + && Objects.equals(prepared.getDatabase(), session.getDatabase())) { + String[] segments = new String[tagValues.length + 1]; + segments[0] = prepared.getTableName(); + System.arraycopy(tagValues, 0, segments, 1, tagValues.length); + deviceId = IDeviceID.Factory.DEFAULT_FACTORY.create(segments); + connection = session.getSessionConnection(deviceId); + } + Long id = prepared.getPreparedId(connection, connection.getSessionId()); + if (id == null) { + TSPrepareNumericRowReq request = new TSPrepareNumericRowReq(); + request + .setTableName(prepared.getTableName()) + .setMeasurements(prepared.getColumnNames()) + .setDataTypes(prepared.getDataTypes().stream().map(TSDataType::ordinal).toList()) + .setColumnCategories( + prepared.getColumnCategories().stream().map(c -> (byte) c.ordinal()).toList()); + // Prepare lazily on each destination; IDs must never cross physical sessions. + id = connection.prepareNumericRow(request); + prepared.bindConnection( + connection, + connection.getSessionId(), + id, + closePreparedRow(connection, connection.getSessionId(), id)); + } int tagBytes = 0; for (String tag : tagValues) { tagBytes += ReadWriteIOUtils.sizeToWrite(tag); @@ -158,19 +188,44 @@ public class TableSession implements ITableSession { } } valueBuffer.flip(); - if (nullBitmap == null) { - connection.executePreparedNumericRow( - prepared.getPreparedId(), timestamp, valueBuffer, tagBuffer); - } else { - connection.executePreparedNumericRow( - prepared.getPreparedId(), - timestamp, - valueBuffer, - tagBuffer, - SessionUtils.encodeNumericRowNullBitmap(nullBitmap, prepared.prepareNullBitmapBuffer())); + try { + if (nullBitmap == null) { + connection.executePreparedNumericRow(id, timestamp, valueBuffer, tagBuffer); + } else { + connection.executePreparedNumericRow( + id, + timestamp, + valueBuffer, + tagBuffer, + SessionUtils.encodeNumericRowNullBitmap( + nullBitmap, prepared.prepareNullBitmapBuffer())); + } + } catch (RedirectException e) { + // This row succeeded. Update only the route for subsequent rows, without replaying this one. + if (deviceId != null && e.getEndPoint() != null) { + session.handleRedirection(deviceId, e.getEndPoint()); + } + } catch (StatementExecutionException e) { + // A server may have invalidated its handle on schema failure. Do not let another cached + // destination silently revive that logical statement; require an explicit new prepare. + prepared.close(); + throw e; } } + private static Runnable closePreparedRow(SessionConnection connection, long sessionId, long id) { + return () -> { + if (connection.getSessionId() != sessionId) { + return; + } + try { + connection.closePreparedNumericRow(id); + } catch (IoTDBConnectionException ignored) { + // Closing the physical session releases any handles left by a failed close RPC. + } + }; + } + @Override public void executeNonQueryStatement(String sql) throws IoTDBConnectionException, StatementExecutionException { diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/PreparedNumericRowRedirectionTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/PreparedNumericRowRedirectionTest.java new file mode 100644 index 00000000000..cdbd984aaf7 --- /dev/null +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/PreparedNumericRowRedirectionTest.java @@ -0,0 +1,219 @@ +/* + * 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.NumericRowPreparedStatement; +import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.rpc.StatementExecutionException; +import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.iotdb.service.rpc.thrift.IClientRPCService; +import org.apache.iotdb.service.rpc.thrift.TSExecutePreparedNumericRowReq; +import org.apache.iotdb.service.rpc.thrift.TSPrepareNumericRowResp; + +import org.apache.thrift.TException; +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.utils.BitMap; +import org.apache.tsfile.utils.ReadWriteIOUtils; +import org.junit.Before; +import org.junit.Test; +import org.mockito.ArgumentCaptor; +import org.powermock.reflect.Whitebox; + +import java.util.HashMap; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class PreparedNumericRowRedirectionTest { + private final TEndPoint leader = new TEndPoint("127.0.0.1", 6668); + private Session session; + private TableSession tableSession; + private SessionConnection initial; + private SessionConnection destination; + private IClientRPCService.Iface initialClient; + private IClientRPCService.Iface destinationClient; + + @Before + public void setUp() throws Exception { + session = new Session.Builder().enableRedirection(true).enableAutoFetch(false).build(); + initialClient = mock(IClientRPCService.Iface.class); + destinationClient = mock(IClientRPCService.Iface.class); + initial = connection(initialClient, 11); + destination = connection(destinationClient, 22); + session.setDefaultSessionConnection(initial); + session.tableModelDeviceIdToEndpoint = new HashMap<>(); + session.endPointToSessionConnection = new HashMap<>(); + session.endPointToSessionConnection.put(leader, destination); + tableSession = new TableSession(session); + when(initialClient.prepareNumericRow(any())) + .thenReturn(new TSPrepareNumericRowResp(success()).setPreparedId(101)); + when(destinationClient.prepareNumericRow(any())) + .thenReturn(new TSPrepareNumericRowResp(success()).setPreparedId(202)); + when(initialClient.executePreparedNumericRow(any())) + .thenReturn( + new TSStatus(TSStatusCode.REDIRECTION_RECOMMEND.getStatusCode()) + .setRedirectNode(leader)); + when(destinationClient.executePreparedNumericRow(any())).thenReturn(success()); + when(initialClient.closePreparedNumericRow(anyLong(), anyLong())).thenReturn(success()); + when(destinationClient.closePreparedNumericRow(anyLong(), anyLong())).thenReturn(success()); + } + + // Recommendations acknowledge this row; later rows prepare once on the leader with intact nulls. + @Test + public void testRedirectPreparesPerSessionWithoutReplaying() throws Exception { + NumericRowPreparedStatement prepared = prepare(); + write(prepared, 1, "a"); + verify(destinationClient, never()).executePreparedNumericRow(any()); + write(prepared, 2, "a"); + write(prepared, 3, "a"); + // A different device has no route yet, so it still uses the original handle. + write(prepared, 4, "b"); + verify(initialClient, times(2)).executePreparedNumericRow(any()); + verify(destinationClient, times(1)).prepareNumericRow(any()); + ArgumentCaptor<TSExecutePreparedNumericRowReq> requests = + ArgumentCaptor.forClass(TSExecutePreparedNumericRowReq.class); + verify(destinationClient, times(2)).executePreparedNumericRow(requests.capture()); + for (TSExecutePreparedNumericRowReq request : requests.getAllValues()) { + assertEquals(22, request.getSessionId()); + assertEquals(202, request.getPreparedId()); + assertEquals(2, request.bufferForNullBitmap().get()); + assertEquals(8, request.bufferForValues().remaining()); + assertEquals(42, request.bufferForValues().getLong()); + assertEquals("a", ReadWriteIOUtils.readString(request.bufferForTagValues())); + } + prepared.close(); + prepared.close(); + verify(initialClient).closePreparedNumericRow(11, 101); + verify(destinationClient).closePreparedNumericRow(22, 202); + } + + // Disabling redirection ignores recommendations; USE must not prepare an old row in a new DB. + @Test + public void testDisabledRedirectionAndDatabaseChange() throws Exception { + try (NumericRowPreparedStatement prepared = prepare()) { + session.setEnableRedirection(false); + write(prepared, 1, "a"); + write(prepared, 2, "a"); + assertTrue(session.tableModelDeviceIdToEndpoint.isEmpty()); + session.setEnableRedirection(true); + write(prepared, 3, "a"); + session.changeDatabase("another_database"); + write(prepared, 4, "a"); + verify(initialClient, times(4)).executePreparedNumericRow(any()); + verify(destinationClient, never()).prepareNumericRow(any()); + } + } + + // Reconnecting a destination requires a new ID; reconnecting the owner invalidates the statement. + @Test + public void testReconnectedSessionsNeverReuseOldIds() throws Exception { + NumericRowPreparedStatement prepared = prepare(); + write(prepared, 1, "a"); + write(prepared, 2, "a"); + Whitebox.setInternalState(destination, "sessionId", 23L); + when(destinationClient.prepareNumericRow(any())) + .thenReturn(new TSPrepareNumericRowResp(success()).setPreparedId(303)); + write(prepared, 3, "a"); + assertEquals(Long.valueOf(303), prepared.getPreparedId(destination, 23)); + verify(destinationClient, times(2)).prepareNumericRow(any()); + Whitebox.setInternalState(initial, "sessionId", 12L); + assertThrows(IoTDBConnectionException.class, () -> write(prepared, 4, "a")); + prepared.close(); + verify(initialClient, never()).closePreparedNumericRow(anyLong(), anyLong()); + verify(destinationClient).closePreparedNumericRow(23, 303); + verify(destinationClient, never()).closePreparedNumericRow(22, 202); + } + + // A failed redirected execution is propagated once, and all handles close on a server error. + @Test + public void testFailureDoesNotReplayOrReviveStatement() throws Exception { + NumericRowPreparedStatement prepared = prepare(); + write(prepared, 1, "a"); + when(destinationClient.executePreparedNumericRow(any())) + .thenThrow(new TException("lost response")); + assertThrows(IoTDBConnectionException.class, () -> write(prepared, 2, "a")); + verify(initialClient).executePreparedNumericRow(any()); + verify(destinationClient).executePreparedNumericRow(any()); + assertFalse(prepared.isClosed()); + doReturn( + new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode()) + .setMessage("schema error")) + .when(destinationClient) + .executePreparedNumericRow(any()); + assertThrows(StatementExecutionException.class, () -> write(prepared, 3, "a")); + assertTrue(prepared.isClosed()); + verify(initialClient).closePreparedNumericRow(11, 101); + verify(destinationClient).closePreparedNumericRow(22, 202); + assertThrows(IoTDBConnectionException.class, () -> write(prepared, 4, "b")); + verify(initialClient).executePreparedNumericRow(any()); + } + + // A preexisting device route is used immediately, including a device with a null TAG. + @Test + public void testExistingRouteAndNullTag() throws Exception { + session.tableModelDeviceIdToEndpoint.put( + IDeviceID.Factory.DEFAULT_FACTORY.create(new String[] {"t", null}), leader); + try (NumericRowPreparedStatement prepared = prepare()) { + write(prepared, 1, null); + verify(initialClient, never()).executePreparedNumericRow(any()); + verify(destinationClient).executePreparedNumericRow(any()); + } + } + + private NumericRowPreparedStatement prepare() throws Exception { + return tableSession.prepareInsertNumericRow( + "t", + List.of("tag", "v", "n"), + List.of(TSDataType.STRING, TSDataType.INT64, TSDataType.INT64), + List.of(ColumnCategory.TAG, ColumnCategory.FIELD, ColumnCategory.FIELD)); + } + + private void write(NumericRowPreparedStatement prepared, long time, String tag) throws Exception { + BitMap nulls = new BitMap(2); + nulls.mark(1); + tableSession.executePreparedInsertNumericRow( + prepared, time, new long[] {42, 99}, new String[] {tag}, nulls); + } + + private static SessionConnection connection(IClientRPCService.Iface client, long sessionId) { + SessionConnection connection = new SessionConnection("table"); + Whitebox.setInternalState(connection, "client", client); + Whitebox.setInternalState(connection, "sessionId", sessionId); + return connection; + } + + private static TSStatus success() { + return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()); + } +}
