Skip to content

Commit 3d7eb21

Browse files
CaideyipiJackieTien97
authored andcommitted
Fix pipe permission retry and table parser progress (#17844)
1 parent 540c008 commit 3d7eb21

3 files changed

Lines changed: 213 additions & 93 deletions

File tree

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/table/TsFileInsertionEventTableParser.java

Lines changed: 108 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -148,39 +148,24 @@ public Iterable<TabletInsertionEvent> toTabletInsertionEvents() {
148148
new Iterator<TabletInsertionEvent>() {
149149

150150
private TsFileInsertionEventTableParserTabletIterator tabletIterator;
151+
private PipeRawTabletInsertionEvent nextEvent;
152+
private Tablet bufferedTablet;
153+
private boolean iterationClosed = false;
151154

152155
@Override
153156
public boolean hasNext() {
154157
try {
155-
if (tabletIterator == null) {
156-
tabletIterator =
157-
new TsFileInsertionEventTableParserTabletIterator(
158-
tsFileSequenceReader,
159-
entry ->
160-
(Objects.isNull(tablePattern)
161-
|| tablePattern.matchesTable(entry.getKey()))
162-
&& hasTablePrivilege(entry.getKey()),
163-
allocatedMemoryBlockForTablet,
164-
allocatedMemoryBlockForBatchData,
165-
allocatedMemoryBlockForChunk,
166-
allocatedMemoryBlockForChunkMeta,
167-
allocatedMemoryBlockForTableSchemas,
168-
currentModifications,
169-
startTime,
170-
endTime);
158+
if (nextEvent != null) {
159+
return true;
171160
}
172-
final boolean hasNext = tabletIterator.hasNext();
173-
if (hasNext && !parseStartTimeRecorded) {
174-
// Record start time on first hasNext() that returns true
175-
recordParseStartTime();
176-
} else if (!hasNext && parseStartTimeRecorded && !parseEndTimeRecorded) {
177-
// Record end time on last hasNext() that returns false
178-
recordParseEndTime();
179-
close();
180-
} else if (!hasNext) {
181-
close();
161+
162+
final Tablet tablet = pollNextNonEmptyTablet();
163+
if (tablet == null) {
164+
return false;
182165
}
183-
return hasNext;
166+
167+
nextEvent = buildTabletInsertionEvent(tablet, !prepareNextNonEmptyTablet());
168+
return true;
184169
} catch (Exception e) {
185170
close();
186171
throw new PipeException(
@@ -211,74 +196,108 @@ private boolean hasTablePrivilege(final String tableName) {
211196
return false;
212197
}
213198

199+
private Tablet pollNextNonEmptyTablet() throws Exception {
200+
if (!prepareNextNonEmptyTablet()) {
201+
return null;
202+
}
203+
204+
final Tablet tablet = bufferedTablet;
205+
bufferedTablet = null;
206+
return tablet;
207+
}
208+
209+
private boolean prepareNextNonEmptyTablet() throws Exception {
210+
if (bufferedTablet != null) {
211+
return true;
212+
}
213+
if (iterationClosed) {
214+
return false;
215+
}
216+
217+
if (tabletIterator == null) {
218+
tabletIterator =
219+
new TsFileInsertionEventTableParserTabletIterator(
220+
tsFileSequenceReader,
221+
entry ->
222+
(Objects.isNull(tablePattern)
223+
|| tablePattern.matchesTable(entry.getKey()))
224+
&& hasTablePrivilege(entry.getKey()),
225+
allocatedMemoryBlockForTablet,
226+
allocatedMemoryBlockForBatchData,
227+
allocatedMemoryBlockForChunk,
228+
allocatedMemoryBlockForChunkMeta,
229+
allocatedMemoryBlockForTableSchemas,
230+
currentModifications,
231+
startTime,
232+
endTime);
233+
}
234+
235+
while (tabletIterator.hasNext()) {
236+
if (!parseStartTimeRecorded) {
237+
recordParseStartTime();
238+
}
239+
240+
final Tablet tablet = tabletIterator.next();
241+
recordTabletMetrics(tablet);
242+
if (!PipeRawTabletInsertionEvent.isTabletEmpty(tablet)) {
243+
bufferedTablet = tablet;
244+
return true;
245+
}
246+
}
247+
248+
closeIteration();
249+
return false;
250+
}
251+
252+
private void closeIteration() {
253+
if (iterationClosed) {
254+
return;
255+
}
256+
257+
if (parseStartTimeRecorded && !parseEndTimeRecorded) {
258+
recordParseEndTime();
259+
}
260+
close();
261+
iterationClosed = true;
262+
}
263+
264+
private PipeRawTabletInsertionEvent buildTabletInsertionEvent(
265+
final Tablet tablet, final boolean needToReport) {
266+
return sourceEvent == null
267+
? new PipeRawTabletInsertionEvent(
268+
Boolean.TRUE,
269+
null,
270+
null,
271+
null,
272+
tablet,
273+
true,
274+
null,
275+
0,
276+
pipeTaskMeta,
277+
sourceEvent,
278+
needToReport)
279+
: new PipeRawTabletInsertionEvent(
280+
Boolean.TRUE,
281+
sourceEvent.getSourceDatabaseNameFromDataRegion(),
282+
sourceEvent.getRawTableModelDataBase(),
283+
sourceEvent.getRawTreeModelDataBase(),
284+
tablet,
285+
true,
286+
sourceEvent.getPipeName(),
287+
sourceEvent.getCreationTime(),
288+
pipeTaskMeta,
289+
sourceEvent,
290+
needToReport);
291+
}
292+
214293
@Override
215294
public TabletInsertionEvent next() {
216295
if (!hasNext()) {
217-
close();
218296
throw new NoSuchElementException();
219297
}
220298

221-
final Tablet tablet = tabletIterator.next();
222-
// Record tablet metrics
223-
recordTabletMetrics(tablet);
224-
225-
final TabletInsertionEvent next;
226-
if (!hasNext()) {
227-
next =
228-
sourceEvent == null
229-
? new PipeRawTabletInsertionEvent(
230-
Boolean.TRUE,
231-
null,
232-
null,
233-
null,
234-
tablet,
235-
true,
236-
null,
237-
0,
238-
pipeTaskMeta,
239-
sourceEvent,
240-
true)
241-
: new PipeRawTabletInsertionEvent(
242-
Boolean.TRUE,
243-
sourceEvent.getSourceDatabaseNameFromDataRegion(),
244-
sourceEvent.getRawTableModelDataBase(),
245-
sourceEvent.getRawTreeModelDataBase(),
246-
tablet,
247-
true,
248-
sourceEvent.getPipeName(),
249-
sourceEvent.getCreationTime(),
250-
pipeTaskMeta,
251-
sourceEvent,
252-
true);
253-
close();
254-
} else {
255-
next =
256-
sourceEvent == null
257-
? new PipeRawTabletInsertionEvent(
258-
Boolean.TRUE,
259-
null,
260-
null,
261-
null,
262-
tablet,
263-
true,
264-
null,
265-
0,
266-
pipeTaskMeta,
267-
sourceEvent,
268-
false)
269-
: new PipeRawTabletInsertionEvent(
270-
Boolean.TRUE,
271-
sourceEvent.getSourceDatabaseNameFromDataRegion(),
272-
sourceEvent.getRawTableModelDataBase(),
273-
sourceEvent.getRawTreeModelDataBase(),
274-
tablet,
275-
true,
276-
sourceEvent.getPipeName(),
277-
sourceEvent.getCreationTime(),
278-
pipeTaskMeta,
279-
sourceEvent,
280-
false);
281-
}
299+
final TabletInsertionEvent next = nextEvent;
300+
nextEvent = null;
282301
return next;
283302
}
284303
};

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/writeback/WriteBackSink.java

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.apache.iotdb.commons.conf.IoTDBConstant;
2525
import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
2626
import org.apache.iotdb.commons.exception.auth.AccessDeniedException;
27+
import org.apache.iotdb.commons.exception.pipe.PipeRuntimeSinkNonReportTimeConfigurableException;
2728
import org.apache.iotdb.commons.pipe.resource.log.PipeLogger;
2829
import org.apache.iotdb.commons.queryengine.common.SqlDialect;
2930
import org.apache.iotdb.commons.utils.StatusUtils;
@@ -284,7 +285,8 @@ private void doTransfer(
284285
&& status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()
285286
&& !(skipIfNoPrivileges
286287
&& status.getCode() == TSStatusCode.NO_PERMISSION.getStatusCode())) {
287-
throw new PipeException(
288+
throwWriteBackExceptionIfNecessary(
289+
status,
288290
String.format(
289291
"Write back PipeInsertNodeTabletInsertionEvent %s error, result status %s",
290292
pipeInsertNodeTabletInsertionEvent, status));
@@ -328,7 +330,8 @@ private void doTransfer(final PipeRawTabletInsertionEvent pipeRawTabletInsertion
328330
&& status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()
329331
&& !(skipIfNoPrivileges
330332
&& status.getCode() == TSStatusCode.NO_PERMISSION.getStatusCode())) {
331-
throw new PipeException(
333+
throwWriteBackExceptionIfNecessary(
334+
status,
332335
String.format(
333336
"Write back PipeRawTabletInsertionEvent %s error, result status %s",
334337
pipeRawTabletInsertionEvent, status));
@@ -373,13 +376,23 @@ private void doTransfer(final PipeStatementInsertionEvent pipeStatementInsertion
373376
&& status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()
374377
&& !(skipIfNoPrivileges
375378
&& status.getCode() == TSStatusCode.NO_PERMISSION.getStatusCode())) {
376-
throw new PipeException(
379+
throwWriteBackExceptionIfNecessary(
380+
status,
377381
String.format(
378382
"Write back PipeStatementInsertionEvent %s error, result status %s",
379383
pipeStatementInsertionEvent, status));
380384
}
381385
}
382386

387+
private static void throwWriteBackExceptionIfNecessary(
388+
final TSStatus status, final String exceptionMessage) {
389+
if (status.getCode() == TSStatusCode.NO_PERMISSION.getStatusCode()) {
390+
throw new PipeRuntimeSinkNonReportTimeConfigurableException(exceptionMessage, Long.MAX_VALUE);
391+
}
392+
393+
throw new PipeException(exceptionMessage);
394+
}
395+
383396
@Override
384397
public void close() throws Exception {
385398
if (session != null) {
@@ -410,7 +423,7 @@ private TSStatus executeStatementForTableModel(
410423
.status;
411424
} catch (final AccessDeniedException e) {
412425
if (!skipIfNoPrivileges) {
413-
throw e;
426+
throw new PipeRuntimeSinkNonReportTimeConfigurableException(e.getMessage(), Long.MAX_VALUE);
414427
}
415428
LOGGER.debug(
416429
DataNodePipeMessages.EXECUTE_STATEMENT_TO_DATABASE_SKIP_BECAUSE_NO,

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionEventParserTest.java

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -525,6 +525,77 @@ public void testTableParserSkipsUnnecessaryBitMaps() throws Exception {
525525
}
526526
}
527527

528+
@Test
529+
public void testTableParserWithTablePatternReportsLastNonEmptyTablet() throws Exception {
530+
final int originalPipeDataStructureTabletRowSize =
531+
PipeConfig.getInstance().getPipeDataStructureTabletRowSize();
532+
CommonDescriptor.getInstance().getConfig().setPipeDataStructureTabletRowSize(2);
533+
534+
try {
535+
alignedTsFile = new File("table-parser-table-pattern.tsfile");
536+
if (alignedTsFile.exists()) {
537+
Assert.assertTrue(alignedTsFile.delete());
538+
}
539+
540+
final List<IMeasurementSchema> schemaList =
541+
Arrays.asList(
542+
new MeasurementSchema("tag0", TSDataType.STRING),
543+
new MeasurementSchema("s0", TSDataType.INT64));
544+
final List<String> columnNameList = Arrays.asList("tag0", "s0");
545+
final List<TSDataType> dataTypeList = Arrays.asList(TSDataType.STRING, TSDataType.INT64);
546+
final List<ColumnCategory> columnCategoryList =
547+
Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD);
548+
549+
try (final TsFileWriter writer = new TsFileWriter(alignedTsFile)) {
550+
writer.registerTableSchema(new TableSchema("test", schemaList, columnCategoryList));
551+
writer.registerTableSchema(new TableSchema("test1", schemaList, columnCategoryList));
552+
writer.writeTable(
553+
generateSimpleTableTablet(
554+
"test", columnNameList, dataTypeList, columnCategoryList, "ignored", 0, 2));
555+
writer.writeTable(
556+
generateSimpleTableTablet(
557+
"test1", columnNameList, dataTypeList, columnCategoryList, "matched", 3, 4));
558+
writer.writeTable(
559+
generateSimpleTableTablet(
560+
"test1", columnNameList, dataTypeList, columnCategoryList, "unmatched", 2, 10));
561+
}
562+
563+
try (final TsFileInsertionEventTableParser parser =
564+
new TsFileInsertionEventTableParser(
565+
alignedTsFile,
566+
new TablePattern(true, null, "test1"),
567+
3,
568+
5,
569+
null,
570+
null,
571+
null,
572+
false)) {
573+
final Iterator<TabletInsertionEvent> iterator = parser.toTabletInsertionEvents().iterator();
574+
int rowCount = 0;
575+
PipeRawTabletInsertionEvent lastEvent = null;
576+
while (iterator.hasNext()) {
577+
final PipeRawTabletInsertionEvent event = (PipeRawTabletInsertionEvent) iterator.next();
578+
final Tablet tablet = event.convertToTablet();
579+
Assert.assertEquals("test1", tablet.getTableName());
580+
Assert.assertFalse(PipeRawTabletInsertionEvent.isTabletEmpty(tablet));
581+
rowCount += tablet.getRowSize();
582+
if (lastEvent != null) {
583+
Assert.assertFalse(lastEvent.isNeedToReport());
584+
}
585+
lastEvent = event;
586+
}
587+
588+
Assert.assertEquals(2, rowCount);
589+
Assert.assertNotNull(lastEvent);
590+
Assert.assertTrue(lastEvent.isNeedToReport());
591+
}
592+
} finally {
593+
CommonDescriptor.getInstance()
594+
.getConfig()
595+
.setPipeDataStructureTabletRowSize(originalPipeDataStructureTabletRowSize);
596+
}
597+
}
598+
528599
@Test
529600
public void manualTestScanParserSplitPerformance() throws Exception {
530601
Assume.assumeTrue(
@@ -1343,6 +1414,23 @@ private void generateLargeAlignedTsFile(
13431414
}
13441415
}
13451416

1417+
private Tablet generateSimpleTableTablet(
1418+
final String tableName,
1419+
final List<String> columnNameList,
1420+
final List<TSDataType> dataTypeList,
1421+
final List<ColumnCategory> columnCategoryList,
1422+
final String tagValue,
1423+
final long... timestamps) {
1424+
final Tablet tablet =
1425+
new Tablet(tableName, columnNameList, dataTypeList, columnCategoryList, timestamps.length);
1426+
for (int rowIndex = 0; rowIndex < timestamps.length; ++rowIndex) {
1427+
tablet.addTimestamp(rowIndex, timestamps[rowIndex]);
1428+
tablet.addValue(rowIndex, 0, tagValue);
1429+
tablet.addValue(rowIndex, 1, (long) rowIndex);
1430+
}
1431+
return tablet;
1432+
}
1433+
13461434
private void generateLargeTableTsFile(
13471435
final File tsFile,
13481436
final int tableCount,

0 commit comments

Comments
 (0)