Skip to content

Commit 3f54368

Browse files
CaideyipiCaideyipi
authored
[to dev/1.3] Fix pipe tree database creation on receiver (#17998)
* Fix pipe tree database creation on receiver (#17991) (cherry picked from commit 28c4e68) * Fix LoadTsFileSchedulerTest import order
1 parent 22bc40e commit 3f54368

5 files changed

Lines changed: 186 additions & 6 deletions

File tree

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222
import org.apache.iotdb.common.rpc.thrift.TSStatus;
2323
import org.apache.iotdb.commons.conf.IoTDBConstant;
2424
import org.apache.iotdb.commons.exception.IllegalPathException;
25+
import org.apache.iotdb.commons.exception.IoTDBException;
26+
import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
2527
import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException;
2628
import org.apache.iotdb.commons.log.LoggerPeriodicalLogReducer;
2729
import org.apache.iotdb.commons.path.PartialPath;
@@ -73,7 +75,9 @@
7375
import org.apache.iotdb.db.queryengine.plan.Coordinator;
7476
import org.apache.iotdb.db.queryengine.plan.analyze.ClusterPartitionFetcher;
7577
import org.apache.iotdb.db.queryengine.plan.analyze.schema.ClusterSchemaFetcher;
78+
import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult;
7679
import org.apache.iotdb.db.queryengine.plan.execution.config.executor.ClusterConfigTaskExecutor;
80+
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.DatabaseSchemaTask;
7781
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metedata.write.view.AlterLogicalViewNode;
7882
import org.apache.iotdb.db.queryengine.plan.statement.Statement;
7983
import org.apache.iotdb.db.queryengine.plan.statement.StatementType;
@@ -96,6 +100,7 @@
96100
import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq;
97101
import org.apache.iotdb.service.rpc.thrift.TPipeTransferResp;
98102

103+
import com.google.common.util.concurrent.ListenableFuture;
99104
import org.apache.tsfile.utils.Pair;
100105
import org.slf4j.Logger;
101106
import org.slf4j.LoggerFactory;
@@ -112,6 +117,8 @@
112117
import java.util.Objects;
113118
import java.util.Optional;
114119
import java.util.Set;
120+
import java.util.concurrent.ConcurrentHashMap;
121+
import java.util.concurrent.ExecutionException;
115122
import java.util.concurrent.atomic.AtomicLong;
116123
import java.util.concurrent.atomic.AtomicReference;
117124
import java.util.stream.Collectors;
@@ -153,6 +160,15 @@ public class IoTDBDataNodeReceiver extends IoTDBFileReceiver {
153160
private static final SessionManager SESSION_MANAGER = SessionManager.getInstance();
154161

155162
private PipeMemoryBlock allocatedMemoryBlock;
163+
private final Set<String> autoCreatedTreeDatabases = ConcurrentHashMap.newKeySet();
164+
private final Set<String> conflictedTreeDatabases = ConcurrentHashMap.newKeySet();
165+
166+
private enum TreeDatabaseCreationResult {
167+
SKIPPED,
168+
CREATED_OR_EXISTED,
169+
CONFLICTED
170+
}
171+
156172
private final List<PipeMemoryBlock> allocatedSliceMemoryBlocks = new ArrayList<>();
157173

158174
static {
@@ -1045,6 +1061,11 @@ private TSStatus executeStatement(final Statement statement) {
10451061
return RpcUtils.getStatus(status.getCode(), status.getMessage());
10461062
}
10471063

1064+
if (autoCreateTreeDatabaseIfNecessary(getTreeDatabaseName(statement))
1065+
== TreeDatabaseCreationResult.CONFLICTED) {
1066+
clearTreeDatabaseName(statement);
1067+
}
1068+
10481069
return Coordinator.getInstance()
10491070
.executeForTreeModel(
10501071
shouldMarkAsPipeRequest.get() ? new PipeEnrichedStatement(statement) : statement,
@@ -1059,6 +1080,91 @@ private TSStatus executeStatement(final Statement statement) {
10591080
.status;
10601081
}
10611082

1083+
private TreeDatabaseCreationResult autoCreateTreeDatabaseIfNecessary(final String database) {
1084+
if (database == null
1085+
|| LoadTsFileStatement.getDatabaseLevelByTreeDatabase(database) == null
1086+
|| !IoTDBDescriptor.getInstance().getConfig().isAutoCreateSchemaEnabled()) {
1087+
return TreeDatabaseCreationResult.SKIPPED;
1088+
}
1089+
if (autoCreatedTreeDatabases.contains(database)) {
1090+
return TreeDatabaseCreationResult.CREATED_OR_EXISTED;
1091+
}
1092+
if (conflictedTreeDatabases.contains(database)) {
1093+
return TreeDatabaseCreationResult.CONFLICTED;
1094+
}
1095+
1096+
try {
1097+
final DatabaseSchemaStatement statement =
1098+
new DatabaseSchemaStatement(DatabaseSchemaStatement.DatabaseSchemaStatementType.CREATE);
1099+
statement.setDatabasePath(new PartialPath(database));
1100+
statement.setEnablePrintExceptionLog(false);
1101+
1102+
final TSStatus permissionStatus = statement.checkPermissionBeforeProcess(username);
1103+
if (permissionStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
1104+
throw new PipeException(permissionStatus.getMessage());
1105+
}
1106+
1107+
final DatabaseSchemaTask task = new DatabaseSchemaTask(statement);
1108+
final ListenableFuture<ConfigTaskResult> future =
1109+
task.execute(ClusterConfigTaskExecutor.getInstance());
1110+
final ConfigTaskResult result = future.get();
1111+
final int statusCode = result.getStatusCode().getStatusCode();
1112+
if (statusCode == TSStatusCode.SUCCESS_STATUS.getStatusCode()
1113+
|| statusCode == TSStatusCode.DATABASE_ALREADY_EXISTS.getStatusCode()) {
1114+
autoCreatedTreeDatabases.add(database);
1115+
return TreeDatabaseCreationResult.CREATED_OR_EXISTED;
1116+
}
1117+
if (statusCode == TSStatusCode.DATABASE_CONFLICT.getStatusCode()) {
1118+
conflictedTreeDatabases.add(database);
1119+
return TreeDatabaseCreationResult.CONFLICTED;
1120+
}
1121+
throw new PipeException(
1122+
String.format(
1123+
"Auto create tree database failed: %s, status: %s",
1124+
database, result.getStatus() == null ? result.getStatusCode() : result.getStatus()));
1125+
} catch (final IllegalPathException e) {
1126+
throw new PipeException(String.format("Illegal tree database %s.", database), e);
1127+
} catch (final ExecutionException | InterruptedException e) {
1128+
if (e instanceof InterruptedException) {
1129+
Thread.currentThread().interrupt();
1130+
}
1131+
final Throwable rootCause = getRootCause(e);
1132+
final int errorCode;
1133+
if (rootCause instanceof IoTDBException) {
1134+
errorCode = ((IoTDBException) rootCause).getErrorCode();
1135+
} else if (rootCause instanceof IoTDBRuntimeException) {
1136+
errorCode = ((IoTDBRuntimeException) rootCause).getErrorCode();
1137+
} else {
1138+
errorCode = TSStatusCode.INTERNAL_SERVER_ERROR.getStatusCode();
1139+
}
1140+
if (errorCode == TSStatusCode.DATABASE_ALREADY_EXISTS.getStatusCode()) {
1141+
autoCreatedTreeDatabases.add(database);
1142+
return TreeDatabaseCreationResult.CREATED_OR_EXISTED;
1143+
}
1144+
if (errorCode == TSStatusCode.DATABASE_CONFLICT.getStatusCode()) {
1145+
conflictedTreeDatabases.add(database);
1146+
return TreeDatabaseCreationResult.CONFLICTED;
1147+
}
1148+
throw new PipeException("Auto create tree database failed because " + e.getMessage(), e);
1149+
}
1150+
}
1151+
1152+
private String getTreeDatabaseName(final Statement statement) {
1153+
if (statement instanceof LoadTsFileStatement) {
1154+
return ((LoadTsFileStatement) statement).getDatabase();
1155+
}
1156+
return null;
1157+
}
1158+
1159+
static void clearTreeDatabaseName(final Statement statement) {
1160+
if (statement instanceof LoadTsFileStatement) {
1161+
final LoadTsFileStatement loadTsFileStatement = (LoadTsFileStatement) statement;
1162+
loadTsFileStatement.setDatabase(null);
1163+
loadTsFileStatement.setDatabaseLevel(
1164+
IoTDBDescriptor.getInstance().getConfig().getDefaultDatabaseLevel());
1165+
}
1166+
}
1167+
10621168
@Override
10631169
protected TSStatus login() {
10641170
final IClientSession session = SESSION_MANAGER.getCurrSession();

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/TreeModelPlanner.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -130,14 +130,19 @@ public IScheduler doSchedule(
130130
&& ((PipeEnrichedStatement) statement).getInnerStatement()
131131
instanceof LoadTsFileStatement;
132132
if (statement instanceof LoadTsFileStatement || isPipeEnrichedTsFileLoad) {
133+
final LoadTsFileStatement loadTsFileStatement =
134+
statement instanceof LoadTsFileStatement
135+
? (LoadTsFileStatement) statement
136+
: (LoadTsFileStatement) ((PipeEnrichedStatement) statement).getInnerStatement();
133137
scheduler =
134138
new LoadTsFileScheduler(
135139
distributedPlan,
136140
context,
137141
stateMachine,
138142
syncInternalServiceClientManager,
139143
partitionFetcher,
140-
isPipeEnrichedTsFileLoad);
144+
isPipeEnrichedTsFileLoad,
145+
loadTsFileStatement.getDatabase());
141146
} else {
142147
scheduler =
143148
new ClusterScheduler(

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@
8282

8383
import java.io.DataOutputStream;
8484
import java.io.File;
85+
import java.io.FileNotFoundException;
8586
import java.io.IOException;
8687
import java.nio.ByteBuffer;
8788
import java.util.ArrayList;
@@ -136,6 +137,7 @@ public class LoadTsFileScheduler implements IScheduler {
136137
private final PlanFragmentId fragmentId;
137138
private final Set<TRegionReplicaSet> allReplicaSets;
138139
private final boolean isGeneratedByPipe;
140+
private final String treeDatabaseForRetry;
139141
private final Map<TTimePartitionSlot, ProgressIndex> timePartitionSlotToProgressIndex;
140142
private final LoadTsFileDataCacheMemoryBlock block;
141143

@@ -145,7 +147,8 @@ public LoadTsFileScheduler(
145147
QueryStateMachine stateMachine,
146148
IClientManager<TEndPoint, SyncDataNodeInternalServiceClient> internalServiceClientManager,
147149
IPartitionFetcher partitionFetcher,
148-
boolean isGeneratedByPipe) {
150+
boolean isGeneratedByPipe,
151+
String treeDatabaseForRetry) {
149152
this.queryContext = queryContext;
150153
this.stateMachine = stateMachine;
151154
this.tsFileNodeList = new ArrayList<>();
@@ -155,6 +158,7 @@ public LoadTsFileScheduler(
155158
this.partitionFetcher = new DataPartitionBatchFetcher(partitionFetcher);
156159
this.allReplicaSets = new HashSet<>();
157160
this.isGeneratedByPipe = isGeneratedByPipe;
161+
this.treeDatabaseForRetry = treeDatabaseForRetry;
158162
this.timePartitionSlotToProgressIndex = new HashMap<>();
159163
this.block = LoadTsFileMemoryManager.getInstance().allocateDataCacheMemoryBlock();
160164

@@ -566,9 +570,7 @@ private void convertFailedTsFilesToTabletsAndRetry() {
566570
final TSStatus status =
567571
loadTsFileDataTypeConverter
568572
.convertForTreeModel(
569-
LoadTsFileStatement.createUnchecked(filePath)
570-
.setDeleteAfterLoad(failedNode.isDeleteAfterLoad())
571-
.setConvertOnTypeMismatch(true))
573+
buildRetryTreeLoadStatement(filePath, failedNode.isDeleteAfterLoad()))
572574
.orElse(null);
573575

574576
if (loadTsFileDataTypeConverter.isSuccessful(status)) {
@@ -607,6 +609,22 @@ private void convertFailedTsFilesToTabletsAndRetry() {
607609
}
608610
}
609611

612+
private LoadTsFileStatement buildRetryTreeLoadStatement(
613+
final String filePath, final boolean deleteAfterLoad) throws FileNotFoundException {
614+
final LoadTsFileStatement statement =
615+
LoadTsFileStatement.createUnchecked(filePath)
616+
.setDeleteAfterLoad(deleteAfterLoad)
617+
.setConvertOnTypeMismatch(true);
618+
if (treeDatabaseForRetry != null) {
619+
statement.setDatabase(treeDatabaseForRetry);
620+
statement.updateDatabaseLevelByTreeDatabase();
621+
}
622+
if (isGeneratedByPipe) {
623+
statement.markIsGeneratedByPipe();
624+
}
625+
return statement;
626+
}
627+
610628
@Override
611629
public void stop(Throwable t) {
612630
dispatcher.abort();

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverTest.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -232,4 +232,23 @@ public void testLoadTsFileSyncStatementCanSkipVerifySchemaWhenNotConvertingType(
232232
Files.deleteIfExists(tsFile);
233233
}
234234
}
235+
236+
@Test
237+
public void testClearTreeDatabaseNameForLoadTsFileStatement() throws Exception {
238+
final Path tsFile = Files.createTempFile("pipe-load-clear-tree-database", ".tsfile");
239+
try {
240+
final LoadTsFileStatement statement =
241+
IoTDBDataNodeReceiver.buildLoadTsFileStatementForSync(
242+
"root.test.sg_0", tsFile.toString(), true, true);
243+
244+
IoTDBDataNodeReceiver.clearTreeDatabaseName(statement);
245+
246+
Assert.assertNull(statement.getDatabase());
247+
Assert.assertEquals(
248+
IoTDBDescriptor.getInstance().getConfig().getDefaultDatabaseLevel(),
249+
statement.getDatabaseLevel());
250+
} finally {
251+
Files.deleteIfExists(tsFile);
252+
}
253+
}
235254
}

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileSchedulerTest.java‎

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment;
2929
import org.apache.iotdb.db.queryengine.plan.planner.plan.SubPlan;
3030
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode;
31+
import org.apache.iotdb.db.queryengine.plan.statement.crud.LoadTsFileStatement;
3132
import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileDataCacheMemoryBlock;
3233

3334
import org.junit.Assert;
@@ -36,9 +37,11 @@
3637
import org.mockito.Mock;
3738
import org.mockito.MockitoAnnotations;
3839

40+
import java.io.File;
3941
import java.lang.reflect.Constructor;
4042
import java.lang.reflect.Field;
4143
import java.lang.reflect.Method;
44+
import java.util.Collections;
4245

4346
import static org.mockito.Mockito.mock;
4447
import static org.mockito.Mockito.spy;
@@ -56,6 +59,7 @@ public void before() {
5659
when(distributedQueryPlan.getRootSubPlan()).thenReturn(subPlan);
5760
when(subPlan.getPlanFragment()).thenReturn(planFragment);
5861
when(planFragment.getId()).thenReturn(new PlanFragmentId("test", 0));
62+
when(distributedQueryPlan.getInstances()).thenReturn(Collections.emptyList());
5963
}
6064

6165
@Test
@@ -68,12 +72,40 @@ public void tt() {
6872
mock(QueryStateMachine.class),
6973
mock(IClientManager.class),
7074
mock(IPartitionFetcher.class),
71-
false));
75+
false,
76+
null));
7277
t.start();
7378
Assert.assertNull(t.getTotalCpuTime());
7479
Assert.assertNull(t.getFragmentInfo());
7580
}
7681

82+
@Test
83+
public void testBuildRetryTreeLoadStatementUpdatesDatabaseLevel() throws Exception {
84+
final LoadTsFileScheduler scheduler =
85+
new LoadTsFileScheduler(
86+
distributedQueryPlan,
87+
mock(MPPQueryContext.class),
88+
mock(QueryStateMachine.class),
89+
mock(IClientManager.class),
90+
mock(IPartitionFetcher.class),
91+
true,
92+
"root.test.sg_0");
93+
final Method method =
94+
LoadTsFileScheduler.class.getDeclaredMethod(
95+
"buildRetryTreeLoadStatement", String.class, boolean.class);
96+
method.setAccessible(true);
97+
98+
final File tsFile = File.createTempFile("test", ".tsfile");
99+
tsFile.deleteOnExit();
100+
101+
final LoadTsFileStatement statement =
102+
(LoadTsFileStatement) method.invoke(scheduler, tsFile.getAbsolutePath(), true);
103+
104+
Assert.assertEquals("root.test.sg_0", statement.getDatabase());
105+
Assert.assertEquals(2, statement.getDatabaseLevel());
106+
Assert.assertTrue(statement.isGeneratedByPipe());
107+
}
108+
77109
@Test
78110
public void testTsFileDataManagerClearReleasesCachedMemory() throws Exception {
79111
final Constructor<LoadTsFileDataCacheMemoryBlock> memoryBlockConstructor =

0 commit comments

Comments
 (0)