diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/PipeTsFileInsertionEvent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/PipeTsFileInsertionEvent.java index 2692000689db..6258dfee1fd8 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/PipeTsFileInsertionEvent.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/PipeTsFileInsertionEvent.java @@ -96,6 +96,16 @@ public class PipeTsFileInsertionEvent extends EnrichedEvent new AtomicReference<>(); private final AtomicReference pendingTabletInsertionEvent = new AtomicReference<>(); + // Only one caller may advance the parser at a time. close() deliberately does not acquire this + // monitor because a parser consumer may be blocked in a user callback. + private final Object tabletConsumptionLock = new Object(); + // Guarded by dataContainer. close() increments this generation to invalidate waiting parsers. + private long parserStateGeneration; + // Guarded by dataContainer. A pending tablet remains owned by its consumer until the callback + // returns, even when close() races with that callback. + private PipeRawTabletInsertionEvent consumingPendingTabletInsertionEvent; + // Guarded by dataContainer. The callback releases the detached tablet after it returns. + private boolean pendingTabletReleaseDeferred; private final AtomicInteger parsedTabletInsertionEventCount = new AtomicInteger(0); private final AtomicBoolean isTsFileParsingCompleted = new AtomicBoolean(false); private final AtomicLong parsedPointCountForCount = new AtomicLong(0); @@ -352,17 +362,33 @@ public boolean internallyIncreaseResourceReferenceCount(final String holderMessa extractTime = System.nanoTime(); final String pipeTsFileResourcePipeName = PipeTsFileResourceManager.getPipeTsFileResourcePipeName(pipeName, creationTime); + final File originalTsFile = tsFile; + final File originalModFile = modFile; + File increasedTsFile = null; + boolean increased = false; try { - tsFile = + increasedTsFile = PipeDataNodeResourceManager.tsfile() .increaseFileReference(tsFile, true, pipeTsFileResourcePipeName); + tsFile = increasedTsFile; if (isWithMod) { modFile = PipeDataNodeResourceManager.tsfile() .increaseFileReference(modFile, false, pipeTsFileResourcePipeName); } + increased = true; return true; } catch (final Exception e) { + if (increasedTsFile != null) { + try { + PipeDataNodeResourceManager.tsfile() + .decreaseFileReference(increasedTsFile, pipeTsFileResourcePipeName); + } catch (final Exception rollbackException) { + e.addSuppressed(rollbackException); + } + } + tsFile = originalTsFile; + modFile = originalModFile; LOGGER.warn( String.format( "Increase reference count for TsFile %s or modFile %s error. Holder Message: %s", @@ -370,9 +396,13 @@ public boolean internallyIncreaseResourceReferenceCount(final String holderMessa e); return false; } finally { - if (Objects.nonNull(pipeName)) { - PipeDataNodeSinglePipeMetrics.getInstance() - .increaseTsFileEventCount(pipeName, creationTime); + if (increased && Objects.nonNull(pipeName)) { + try { + PipeDataNodeSinglePipeMetrics.getInstance() + .increaseTsFileEventCount(pipeName, creationTime); + } catch (final Exception e) { + LOGGER.warn("Decrease reference count for TsFile {} error.", tsFile, e); + } } } } @@ -381,31 +411,57 @@ public boolean internallyIncreaseResourceReferenceCount(final String holderMessa public boolean internallyDecreaseResourceReferenceCount(final String holderMessage) { final String pipeTsFileResourcePipeName = PipeTsFileResourceManager.getPipeTsFileResourcePipeName(pipeName, creationTime); + boolean isSuccessful = true; try { PipeDataNodeResourceManager.tsfile() .decreaseFileReference(tsFile, pipeTsFileResourcePipeName); - if (isWithMod) { + } catch (final Exception e) { + isSuccessful = false; + LOGGER.warn( + String.format( + "Decrease reference count for TsFile %s error. Holder Message: %s", + tsFile, holderMessage), + e); + } + + if (isWithMod) { + try { PipeDataNodeResourceManager.tsfile() .decreaseFileReference(modFile, pipeTsFileResourcePipeName); + } catch (final Exception e) { + isSuccessful = false; + LOGGER.warn( + String.format( + "Decrease reference count for TsFile %s error. Holder Message: %s", + modFile, holderMessage), + e); } + } + + try { close(); - return true; } catch (final Exception e) { + isSuccessful = false; LOGGER.warn( String.format( "Decrease reference count for TsFile %s error. Holder Message: %s", - tsFile.getPath(), holderMessage), + tsFile, holderMessage), e); - return false; } finally { if (Objects.nonNull(pipeName)) { - PipeDataNodeSinglePipeMetrics.getInstance() - .decreaseTsFileEventCount( - pipeName, - creationTime, - shouldReportOnCommit ? System.nanoTime() - extractTime : -1); + try { + PipeDataNodeSinglePipeMetrics.getInstance() + .decreaseTsFileEventCount( + pipeName, + creationTime, + shouldReportOnCommit ? System.nanoTime() - extractTime : -1); + } catch (final Exception e) { + isSuccessful = false; + LOGGER.warn("Decrease reference count for TsFile {} error.", tsFile, e); + } } } + return isSuccessful; } @Override @@ -567,83 +623,157 @@ public void consumeTabletInsertionEventsWithRetry( final String callerName, final PipeProcessorSubtaskExecutionGuard processorExecutionGuard) throws Exception { - try { - while (true) { - processorExecutionGuard.check(); - final PipeRawTabletInsertionEvent parsedEvent = - getNextTabletInsertionEventFromSavedProgress(processorExecutionGuard); - if (parsedEvent == null) { - isTsFileParsingCompleted.set(true); - releaseTsFileParserMemoryIfReserved(); - return; + synchronized (tabletConsumptionLock) { + final long parserStateGeneration; + synchronized (dataContainer) { + parserStateGeneration = this.parserStateGeneration; + } + + try { + while (true) { + processorExecutionGuard.check(); + final PipeRawTabletInsertionEvent parsedEvent = + getNextTabletInsertionEventFromSavedProgress( + processorExecutionGuard, parserStateGeneration); + if (parsedEvent == null) { + synchronized (dataContainer) { + if (parserStateGeneration == this.parserStateGeneration) { + isTsFileParsingCompleted.set(true); + } + } + releaseTsFileParserMemoryIfReserved(); + return; + } + + boolean consumed = false; + try { + processorExecutionGuard.check(); + consumeParsedTabletInsertionEventWithRetry( + consumer, + callerName, + parsedTabletInsertionEventCount.get(), + parsedEvent, + processorExecutionGuard); + consumed = true; + } finally { + finishConsumingTabletInsertionEvent(parsedEvent, consumed); + } + + synchronized (dataContainer) { + if (parserStateGeneration != this.parserStateGeneration) { + return; + } + } + processorExecutionGuard.check(); } - processorExecutionGuard.check(); - consumeParsedTabletInsertionEventWithRetry( - consumer, + } catch (final PipeProcessorSubtaskYieldException e) { + releaseTsFileParserMemoryIfReserved(); + if (!processorExecutionGuard.isCurrentInvocationValid()) { + cancelTsFileParserMemoryReservationIfPending(); + } + throw e; + } catch (final PipeRuntimeOutOfMemoryCriticalException e) { + // Yield the active parser slot to the next pipe while retaining the iterator and current + // tablet. The next retry resumes from this exact tablet instead of reparsing the TsFile. + releaseTsFileParserMemoryIfReserved(); + LOGGER.warn( + "{}: failed to allocate memory for parsing TsFile {}, tablet event no. {}, will release parser memory and retry the TsFile event later.", callerName, + getTsFile(), parsedTabletInsertionEventCount.get(), - parsedEvent, - processorExecutionGuard); - pendingTabletInsertionEvent.compareAndSet(parsedEvent, null); - processorExecutionGuard.check(); - } - } catch (final PipeProcessorSubtaskYieldException e) { - releaseTsFileParserMemoryIfReserved(); - if (!processorExecutionGuard.isCurrentInvocationValid()) { - cancelTsFileParserMemoryReservationIfPending(); + e); + throw e; + } catch (final Exception e) { + releaseTsFileParserMemoryIfReserved(); + throw e; } - throw e; - } catch (final PipeRuntimeOutOfMemoryCriticalException e) { - // Yield the active parser slot to the next pipe while retaining the iterator and current - // tablet. The next retry resumes from this exact tablet instead of reparsing the TsFile. - releaseTsFileParserMemoryIfReserved(); - LOGGER.warn( - "{}: failed to allocate memory for parsing TsFile {}, tablet event no. {}, will release parser memory and retry the TsFile event later.", - callerName, - getTsFile(), - parsedTabletInsertionEventCount.get(), - e); - throw e; - } catch (final Exception e) { - releaseTsFileParserMemoryIfReserved(); - throw e; } } private PipeRawTabletInsertionEvent getNextTabletInsertionEventFromSavedProgress( - final PipeProcessorSubtaskExecutionGuard processorExecutionGuard) throws Exception { - if (isTsFileParsingCompleted.get()) { - return null; + final PipeProcessorSubtaskExecutionGuard processorExecutionGuard, + final long expectedParserStateGeneration) + throws Exception { + synchronized (dataContainer) { + if (expectedParserStateGeneration != parserStateGeneration + || isTsFileParsingCompleted.get()) { + return null; + } } // Reacquire parser memory after a previous failure yielded the active parser slot. Processor // subtasks use non-blocking admission here, while other callers retain the bounded wait. - reserveResource4Parsing(processorExecutionGuard); - - final PipeRawTabletInsertionEvent pendingEvent = pendingTabletInsertionEvent.get(); - if (pendingEvent != null) { - return pendingEvent; + if (!reserveResource4Parsing(processorExecutionGuard, expectedParserStateGeneration)) { + return null; } - Iterator iterator = tabletInsertionEventIterator.get(); - if (iterator == null) { - if (!waitForTsFileClose(processorExecutionGuard)) { - LOGGER.warn( - "Pipe skipping temporary TsFile's parsing which shouldn't be transferred: {}", tsFile); + synchronized (dataContainer) { + if (expectedParserStateGeneration != parserStateGeneration + || isTsFileParsingCompleted.get()) { return null; } - iterator = initDataContainer().toTabletInsertionEvents().iterator(); - tabletInsertionEventIterator.set(iterator); + + final PipeRawTabletInsertionEvent pendingEvent = pendingTabletInsertionEvent.get(); + if (pendingEvent != null) { + consumingPendingTabletInsertionEvent = pendingEvent; + return pendingEvent; + } } - if (!iterator.hasNext()) { + if (!waitForTsFileClose(processorExecutionGuard)) { + LOGGER.warn( + "Pipe skipping temporary TsFile's parsing which shouldn't be transferred: {}", tsFile); return null; } - final PipeRawTabletInsertionEvent nextEvent = (PipeRawTabletInsertionEvent) iterator.next(); - pendingTabletInsertionEvent.set(nextEvent); - parsedTabletInsertionEventCount.incrementAndGet(); - return nextEvent; + synchronized (dataContainer) { + if (expectedParserStateGeneration != parserStateGeneration + || isTsFileParsingCompleted.get()) { + return null; + } + + Iterator iterator = tabletInsertionEventIterator.get(); + if (iterator == null) { + iterator = initDataContainer().toTabletInsertionEvents().iterator(); + if (expectedParserStateGeneration != parserStateGeneration) { + return null; + } + tabletInsertionEventIterator.set(iterator); + } + + if (!iterator.hasNext()) { + return null; + } + + final PipeRawTabletInsertionEvent nextEvent = (PipeRawTabletInsertionEvent) iterator.next(); + pendingTabletInsertionEvent.set(nextEvent); + parsedTabletInsertionEventCount.incrementAndGet(); + consumingPendingTabletInsertionEvent = nextEvent; + return nextEvent; + } + } + + private void finishConsumingTabletInsertionEvent( + final PipeRawTabletInsertionEvent event, final boolean consumed) { + PipeRawTabletInsertionEvent eventToRelease = null; + synchronized (dataContainer) { + if (consumingPendingTabletInsertionEvent != event) { + return; + } + + consumingPendingTabletInsertionEvent = null; + if (consumed && pendingTabletInsertionEvent.get() == event) { + pendingTabletInsertionEvent.compareAndSet(event, null); + } + if (pendingTabletReleaseDeferred) { + pendingTabletReleaseDeferred = false; + eventToRelease = event; + } + } + + if (eventToRelease != null) { + releaseParsedTabletEvent(eventToRelease); + } } private void consumeParsedTabletInsertionEventWithRetry( @@ -753,8 +883,19 @@ public Iterable toTabletInsertionEvents(final long timeout "Pipe skipping temporary TsFile's parsing which shouldn't be transferred: {}", tsFile); return Collections.emptyList(); } - waitForResourceEnough4Parsing(timeoutMs); - return initDataContainer().toTabletInsertionEvents(); + final long parserStateGeneration; + synchronized (dataContainer) { + parserStateGeneration = this.parserStateGeneration; + } + if (!waitForResourceEnough4Parsing(timeoutMs, parserStateGeneration)) { + return Collections.emptyList(); + } + synchronized (dataContainer) { + if (parserStateGeneration != this.parserStateGeneration) { + return Collections.emptyList(); + } + return initDataContainer().toTabletInsertionEvents(); + } } catch (final Exception e) { close(); @@ -778,20 +919,28 @@ public Iterable toTabletInsertionEvents(final long timeout } } - private void reserveResource4Parsing( - final PipeProcessorSubtaskExecutionGuard processorExecutionGuard) + private boolean reserveResource4Parsing( + final PipeProcessorSubtaskExecutionGuard processorExecutionGuard, + final long expectedParserStateGeneration) throws InterruptedException { if (!processorExecutionGuard.isEnabled()) { - waitForResourceEnough4Parsing((long) ((1 + Math.random()) * 20 * 1000)); - return; + return waitForResourceEnough4Parsing( + (long) ((1 + Math.random()) * 20 * 1000), expectedParserStateGeneration); } processorExecutionGuard.check(); + if (!isParserStateGenerationCurrent(expectedParserStateGeneration)) { + return false; + } final PipeMemoryManager memoryManager = PipeDataNodeResourceManager.memory(); if (tryReserveTsFileParserMemory(memoryManager)) { try { processorExecutionGuard.check(); - return; + if (isParserStateGenerationCurrent(expectedParserStateGeneration)) { + return true; + } + releaseTsFileParserMemoryIfReserved(); + return false; } catch (final PipeProcessorSubtaskYieldException e) { releaseTsFileParserMemoryIfReserved(); throw e; @@ -802,19 +951,36 @@ private void reserveResource4Parsing( cancelTsFileParserMemoryReservationIfPending(); processorExecutionGuard.check(); } + if (!isParserStateGenerationCurrent(expectedParserStateGeneration)) { + cancelTsFileParserMemoryReservationIfPending(); + return false; + } processorExecutionGuard.yieldIfParserNotAdmitted(); + return false; } - private void waitForResourceEnough4Parsing(final long timeoutMs) throws InterruptedException { + private boolean waitForResourceEnough4Parsing( + final long timeoutMs, final long expectedParserStateGeneration) throws InterruptedException { final PipeMemoryManager memoryManager = PipeDataNodeResourceManager.memory(); + if (!isParserStateGenerationCurrent(expectedParserStateGeneration)) { + return false; + } if (tryReserveTsFileParserMemory(memoryManager)) { - return; + if (isParserStateGenerationCurrent(expectedParserStateGeneration)) { + return true; + } + releaseTsFileParserMemoryIfReserved(); + return false; } final long startTime = System.currentTimeMillis(); long lastRecordTime = startTime; while (!tryReserveTsFileParserMemory(memoryManager)) { + if (!isParserStateGenerationCurrent(expectedParserStateGeneration)) { + cancelTsFileParserMemoryReservationIfPending(); + return false; + } final long currentTime = System.currentTimeMillis(); final long elapsedRecordTimeInMs = currentTime - lastRecordTime; final long waitTimeInMs = currentTime - startTime; @@ -852,6 +1018,17 @@ private void waitForResourceEnough4Parsing(final long timeoutMs) throws Interrup "Wait for memory enough for parsing {} for {} seconds.", resource != null ? resource.getTsFilePath() : "tsfile", waitTimeSeconds); + if (isParserStateGenerationCurrent(expectedParserStateGeneration)) { + return true; + } + releaseTsFileParserMemoryIfReserved(); + return false; + } + + private boolean isParserStateGenerationCurrent(final long expectedParserStateGeneration) { + synchronized (dataContainer) { + return expectedParserStateGeneration == parserStateGeneration; + } } private boolean tryReserveTsFileParserMemory(final PipeMemoryManager memoryManager) { @@ -880,10 +1057,12 @@ private void releaseTsFileParserMemoryIfReserved() { } public void cancelTsFileParserMemoryReservationIfPending() { - if (!isTsFileParserMemoryReserved.get()) { - PipeDataNodeResourceManager.memory() - .cancelTsFileParserMemoryReservation( - pipeName, creationTime, dataRegionId, tsFileParserMemoryReservationKey); + synchronized (isTsFileParserMemoryReserved) { + if (!isTsFileParserMemoryReserved.get()) { + PipeDataNodeResourceManager.memory() + .cancelTsFileParserMemoryReservation( + pipeName, creationTime, dataRegionId, tsFileParserMemoryReservationKey); + } } } @@ -898,20 +1077,26 @@ public boolean isGeneratedByHistoricalExtractor() { private TsFileInsertionDataContainer initDataContainer() { try { - dataContainer.compareAndSet( - null, - new TsFileInsertionDataContainerProvider( - pipeName, - creationTime, - tsFile, - pipePattern, - startTime, - endTime, - pipeTaskMeta, - this, - tsFileParser) - .provide(isWithMod)); - return dataContainer.get(); + synchronized (dataContainer) { + final TsFileInsertionDataContainer container = dataContainer.get(); + if (container != null) { + return container; + } + final TsFileInsertionDataContainer createdContainer = + new TsFileInsertionDataContainerProvider( + pipeName, + creationTime, + tsFile, + pipePattern, + startTime, + endTime, + pipeTaskMeta, + this, + tsFileParser) + .provide(isWithMod); + dataContainer.set(createdContainer); + return createdContainer; + } } catch (final IOException e) { close(); @@ -949,20 +1134,53 @@ public long count(final boolean skipReportOnCommit) throws Exception { /** Release the resource of {@link TsFileInsertionDataContainer}. */ @Override public void close() { - cancelTsFileParserMemoryReservationIfPending(); - tabletInsertionEventIterator.set(null); - releaseParsedTabletEvent(pendingTabletInsertionEvent.getAndSet(null)); - parsedTabletInsertionEventCount.set(0); - parsedPointCountForCount.set(0); - isTsFileParsingCompleted.set(false); - dataContainer.getAndUpdate( - container -> { - if (Objects.nonNull(container)) { - container.close(); - } - return null; - }); - releaseTsFileParserMemoryIfReserved(); + try { + cancelTsFileParserMemoryReservationIfPending(); + } catch (final Exception e) { + LOGGER.warn("Decrease reference count for TsFile {} error.", tsFile, e); + } + + final PipeRawTabletInsertionEvent detachedPendingEvent; + final boolean pendingEventIsBeingConsumed; + final TsFileInsertionDataContainer containerToClose; + synchronized (dataContainer) { + ++parserStateGeneration; + tabletInsertionEventIterator.set(null); + detachedPendingEvent = pendingTabletInsertionEvent.getAndSet(null); + parsedTabletInsertionEventCount.set(0); + parsedPointCountForCount.set(0); + isTsFileParsingCompleted.set(false); + + containerToClose = dataContainer.getAndSet(null); + pendingEventIsBeingConsumed = + detachedPendingEvent != null + && detachedPendingEvent == consumingPendingTabletInsertionEvent; + if (pendingEventIsBeingConsumed) { + pendingTabletReleaseDeferred = true; + } + } + + if (detachedPendingEvent != null && !pendingEventIsBeingConsumed) { + try { + releaseParsedTabletEvent(detachedPendingEvent); + } catch (final Exception e) { + LOGGER.warn("Decrease reference count for TsFile {} error.", tsFile, e); + } + } + + if (containerToClose != null) { + try { + containerToClose.close(); + } catch (final Exception e) { + LOGGER.warn("Decrease reference count for TsFile {} error.", tsFile, e); + } + } + + try { + releaseTsFileParserMemoryIfReserved(); + } catch (final Exception e) { + LOGGER.warn("Decrease reference count for TsFile {} error.", tsFile, e); + } } /////////////////////////// Object /////////////////////////// @@ -1050,24 +1268,37 @@ protected void finalizeResource() { PipeDataNodeResourceManager.memory() .cancelTsFileParserMemoryReservation( pipeName, creationTime, dataRegionId, tsFileParserMemoryReservationKey); - final String pipeTsFileResourcePipeName = - PipeTsFileResourceManager.getPipeTsFileResourcePipeName(pipeName, creationTime); - // decrease reference count + } catch (final Exception e) { + LOGGER.warn("Decrease reference count for TsFile {} error.", tsFile, e); + } + + final String pipeTsFileResourcePipeName = + PipeTsFileResourceManager.getPipeTsFileResourcePipeName(pipeName, creationTime); + try { PipeDataNodeResourceManager.tsfile() .decreaseFileReference(tsFile, pipeTsFileResourcePipeName); - if (isWithMod) { + } catch (final Exception e) { + LOGGER.warn("Decrease reference count for TsFile {} error.", tsFile, e); + } + if (isWithMod) { + try { PipeDataNodeResourceManager.tsfile() .decreaseFileReference(modFile, pipeTsFileResourcePipeName); + } catch (final Exception e) { + LOGGER.warn("Decrease reference count for TsFile {} error.", modFile, e); } + } - // close data container - dataContainer.getAndUpdate( - container -> { - if (Objects.nonNull(container)) { - container.close(); - } - return null; - }); + try { + final TsFileInsertionDataContainer container = dataContainer.getAndSet(null); + if (container != null) { + container.close(); + } + } catch (final Exception e) { + LOGGER.warn("Decrease reference count for TsFile {} error.", tsFile, e); + } + + try { synchronized (isTsFileParserMemoryReserved) { if (isTsFileParserMemoryReserved.compareAndSet(true, false)) { PipeDataNodeResourceManager.memory() @@ -1075,7 +1306,7 @@ protected void finalizeResource() { } } } catch (final Exception e) { - LOGGER.warn("Decrease reference count for TsFile {} error.", tsFile.getPath(), e); + LOGGER.warn("Decrease reference count for TsFile {} error.", tsFile, e); } } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryManager.java index 68669ce4b55d..abb1dfebe41e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryManager.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.pipe.event.common.tsfile.parser; +import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager; import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryBlock; @@ -64,7 +65,17 @@ public long getMemoryUsageInBytes() { @Override public void forceResize(final long newSizeInBytes) { - PipeDataNodeResourceManager.memory().forceResize(delegate, newSizeInBytes); + final long oldSize = delegate.getMemoryUsageInBytes(); + if (!PipeDataNodeResourceManager.memory().tryResize(delegate, newSizeInBytes)) { + throw new PipeRuntimeOutOfMemoryCriticalException( + String.format( + "forceResize: failed to allocate memory after %d retries, total memory size %d " + + "bytes, used memory size %d bytes, requested memory size %d bytes", + 0, + PipeDataNodeResourceManager.memory().getTotalNonFloatingMemorySizeInBytes(), + PipeDataNodeResourceManager.memory().getUsedMemorySizeInBytes(), + newSizeInBytes - oldSize)); + } } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java index 9c36c8dbf7ea..5f03991d3501 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java @@ -690,6 +690,24 @@ public void forceResize(final PipeMemoryBlock block, final long targetSize) { resize(block, targetSize, true); } + /** + * Attempts a single resize without waiting for other pipe tasks to release memory. + * + *

This is intended for callers that hold payload/batch locks and can actively release memory + * after a failed attempt. Waiting in that situation can prevent the caller itself from making + * forward progress. + */ + public synchronized boolean tryResize(final PipeMemoryBlock block, final long targetSize) { + if (targetSize < 0) { + return false; + } + if (block == null || block.isReleased()) { + LOGGER.warn("forceResize: cannot resize a null or released memory block"); + return false; + } + return tryResizeInternal(block, targetSize); + } + public synchronized void resize( final PipeMemoryBlock block, final long targetSize, final boolean force) { if (block == null || block.isReleased()) { @@ -697,59 +715,18 @@ public synchronized void resize( return; } - if (!PIPE_MEMORY_MANAGEMENT_ENABLED) { - block.setMemoryUsageInBytes(targetSize); + if (tryResizeInternal(block, targetSize)) { return; } - final long oldSize = block.getMemoryUsageInBytes(); - if (oldSize >= targetSize) { - usedMemorySizeInBytes -= oldSize - targetSize; - if (block instanceof PipeTabletMemoryBlock) { - usedMemorySizeInBytesOfTablets -= oldSize - targetSize; - } - if (block instanceof PipeTsFileMemoryBlock) { - usedMemorySizeInBytesOfTsFiles -= oldSize - targetSize; - } - block.setMemoryUsageInBytes(targetSize); - - // If no memory is used in the block, we can remove it from the allocated blocks. - if (targetSize == 0) { - allocatedBlocks.remove(block); - } - - notifyNextTsFileParserMemoryReservationInternal(); - this.notifyAll(); - return; - } - - long sizeInBytes = targetSize - oldSize; + final long sizeInBytes = targetSize - block.getMemoryUsageInBytes(); final int memoryAllocateMaxRetries = PipeConfig.getInstance().getPipeMemoryAllocateMaxRetries(); for (int i = 1; i <= memoryAllocateMaxRetries; i++) { - // Dynamically resized data-structure blocks must obey the same admission thresholds as - // blocks allocated with a non-zero initial size. Otherwise they can exhaust the pool and - // prevent downstream consumers from allocating the memory needed to release them. - if (isHardEnoughForResizing(block, sizeInBytes) - && getTotalNonFloatingMemorySizeInBytes() - usedMemorySizeInBytes >= sizeInBytes) { - usedMemorySizeInBytes += sizeInBytes; - if (oldSize == 0) { - // If the memory block is not registered, we need to register it first. - // Otherwise, the memory usage will be inconsistent. - // See registerMemoryBlock for more details. - allocatedBlocks.add(block); - } - if (block instanceof PipeTabletMemoryBlock) { - usedMemorySizeInBytesOfTablets += sizeInBytes; - } - if (block instanceof PipeTsFileMemoryBlock) { - usedMemorySizeInBytesOfTsFiles += sizeInBytes; - } - block.setMemoryUsageInBytes(targetSize); - return; - } - try { tryShrinkUntilFreeMemorySatisfy(sizeInBytes); + if (tryResizeInternal(block, targetSize)) { + return; + } this.wait(PipeConfig.getInstance().getPipeMemoryAllocateRetryIntervalInMs()); } catch (InterruptedException e) { Thread.currentThread().interrupt(); @@ -770,6 +747,57 @@ && getTotalNonFloatingMemorySizeInBytes() - usedMemorySizeInBytes >= sizeInBytes } } + private boolean tryResizeInternal(final PipeMemoryBlock block, final long targetSize) { + if (!PIPE_MEMORY_MANAGEMENT_ENABLED) { + block.setMemoryUsageInBytes(targetSize); + return true; + } + + final long oldSize = block.getMemoryUsageInBytes(); + if (oldSize >= targetSize) { + final long releasedSize = oldSize - targetSize; + usedMemorySizeInBytes -= releasedSize; + if (block instanceof PipeTabletMemoryBlock) { + usedMemorySizeInBytesOfTablets -= releasedSize; + } + if (block instanceof PipeTsFileMemoryBlock) { + usedMemorySizeInBytesOfTsFiles -= releasedSize; + } + block.setMemoryUsageInBytes(targetSize); + + if (targetSize == 0) { + allocatedBlocks.remove(block); + } + + notifyNextTsFileParserMemoryReservationInternal(); + this.notifyAll(); + return true; + } + + final long sizeInBytes = targetSize - oldSize; + // Dynamically resized data-structure blocks must obey the same admission thresholds as blocks + // allocated with a non-zero initial size. Otherwise they can exhaust the pool and prevent + // downstream consumers from allocating the memory needed to release them. + if (!isHardEnoughForResizing(block, sizeInBytes) + || getTotalNonFloatingMemorySizeInBytes() - usedMemorySizeInBytes < sizeInBytes) { + return false; + } + + usedMemorySizeInBytes += sizeInBytes; + if (oldSize == 0) { + // Zero-sized blocks are registered lazily on their first successful expansion. + allocatedBlocks.add(block); + } + if (block instanceof PipeTabletMemoryBlock) { + usedMemorySizeInBytesOfTablets += sizeInBytes; + } + if (block instanceof PipeTsFileMemoryBlock) { + usedMemorySizeInBytesOfTsFiles += sizeInBytes; + } + block.setMemoryUsageInBytes(targetSize); + return true; + } + /** * Allocate a {@link PipeMemoryBlock} for pipe only if memory used after allocation is less than * the specified threshold. diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFilePublicResource.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFilePublicResource.java index 47134fe117c9..48b2b96c8085 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFilePublicResource.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFilePublicResource.java @@ -37,7 +37,6 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Objects; public class PipeTsFilePublicResource extends PipeTsFileResource { private static final Logger LOGGER = LoggerFactory.getLogger(PipeTsFilePublicResource.class); @@ -52,24 +51,15 @@ public PipeTsFilePublicResource(File hardlinkOrCopiedFile) { } @Override - public void close() { + public synchronized void close() { super.close(); - - if (deviceMeasurementsMap != null) { - deviceMeasurementsMap = null; - } - - if (deviceIsAlignedMap != null) { - deviceIsAlignedMap = null; - } - - if (measurementDataTypeMap != null) { - measurementDataTypeMap = null; - } - - if (allocatedMemoryBlock != null) { - allocatedMemoryBlock.close(); - allocatedMemoryBlock = null; + deviceMeasurementsMap = null; + deviceIsAlignedMap = null; + measurementDataTypeMap = null; + final PipeMemoryBlock block = allocatedMemoryBlock; + allocatedMemoryBlock = null; + if (block != null) { + block.close(); } } @@ -113,46 +103,53 @@ synchronized boolean cacheDeviceIsAlignedMapIfAbsent(final File tsFile) throws I // See if pipe memory is sufficient to be allocated for TsFileSequenceReader. // Only allocate when pipe memory used is less than 50%, because memory here // is hard to shrink and may consume too much memory. - allocatedMemoryBlock = + final PipeMemoryBlock readerMemoryBlock = PipeDataNodeResourceManager.memory() .forceAllocateIfSufficient( PipeConfig.getInstance().getPipeMemoryAllocateForTsFileSequenceReaderInBytes(), MEMORY_SUFFICIENT_THRESHOLD); - if (allocatedMemoryBlock == null) { + if (readerMemoryBlock == null) { LOGGER.info( "Failed to cacheDeviceIsAlignedMapIfAbsent for tsfile {}, because memory usage is high", tsFile.getPath()); return false; } + final Map cachedDeviceIsAlignedMap = new HashMap<>(); long memoryRequiredInBytes = 0L; - try (TsFileSequenceReader sequenceReader = - new TsFileSequenceReader(tsFile.getPath(), true, false)) { - deviceIsAlignedMap = new HashMap<>(); - final TsFileDeviceIterator deviceIsAlignedIterator = - sequenceReader.getAllDevicesIteratorWithIsAligned(); - while (deviceIsAlignedIterator.hasNext()) { - final Pair deviceIsAlignedPair = deviceIsAlignedIterator.next(); - deviceIsAlignedMap.put(deviceIsAlignedPair.getLeft(), deviceIsAlignedPair.getRight()); + try { + try (TsFileSequenceReader sequenceReader = + new TsFileSequenceReader(tsFile.getPath(), true, false)) { + final TsFileDeviceIterator deviceIsAlignedIterator = + sequenceReader.getAllDevicesIteratorWithIsAligned(); + while (deviceIsAlignedIterator.hasNext()) { + final Pair deviceIsAlignedPair = deviceIsAlignedIterator.next(); + cachedDeviceIsAlignedMap.put( + deviceIsAlignedPair.getLeft(), deviceIsAlignedPair.getRight()); + } } - memoryRequiredInBytes += PipeMemoryWeightUtil.memoryOfIDeviceId2Bool(deviceIsAlignedMap); + memoryRequiredInBytes += + PipeMemoryWeightUtil.memoryOfIDeviceId2Bool(cachedDeviceIsAlignedMap); + } finally { + // The reader block is temporary and must never become the persistent metadata block. + readerMemoryBlock.close(); } - // Release memory of TsFileSequenceReader. - allocatedMemoryBlock.close(); - allocatedMemoryBlock = null; // Allocate again for the cached objects. - allocatedMemoryBlock = + final PipeMemoryBlock cachedMemoryBlock = PipeDataNodeResourceManager.memory() .forceAllocateIfSufficient(memoryRequiredInBytes, MEMORY_SUFFICIENT_THRESHOLD); - if (allocatedMemoryBlock == null) { + if (cachedMemoryBlock == null) { LOGGER.info( "PipeTsFileResource: Failed to cache objects for tsfile {} in cache, because memory usage is high", tsFile.getPath()); - deviceIsAlignedMap = null; return false; } + // Publish the map only after its accounting block has been acquired. Readers never observe + // a partially built map or a map without a corresponding memory reservation. + deviceIsAlignedMap = cachedDeviceIsAlignedMap; + allocatedMemoryBlock = cachedMemoryBlock; LOGGER.info("PipeTsFileResource: Cached deviceIsAlignedMap for tsfile {}.", tsFile.getPath()); return true; } @@ -163,65 +160,75 @@ synchronized boolean cacheObjectsIfAbsent(final File tsFile) throws IOException return true; } else { // Recalculate it again because only deviceIsAligned map is cached - allocatedMemoryBlock.close(); + final PipeMemoryBlock oldMemoryBlock = allocatedMemoryBlock; allocatedMemoryBlock = null; + deviceIsAlignedMap = null; + oldMemoryBlock.close(); } } // See if pipe memory is sufficient to be allocated for TsFileSequenceReader. // Only allocate when pipe memory used is less than 50%, because memory here // is hard to shrink and may consume too much memory. - allocatedMemoryBlock = + final PipeMemoryBlock readerMemoryBlock = PipeDataNodeResourceManager.memory() .forceAllocateIfSufficient( PipeConfig.getInstance().getPipeMemoryAllocateForTsFileSequenceReaderInBytes(), MEMORY_SUFFICIENT_THRESHOLD); - if (allocatedMemoryBlock == null) { + if (readerMemoryBlock == null) { LOGGER.info( "Failed to cacheObjectsIfAbsent for tsfile {}, because memory usage is high", tsFile.getPath()); return false; } + Map> cachedDeviceMeasurementsMap = null; + Map cachedDeviceIsAlignedMap = null; + Map cachedMeasurementDataTypeMap = null; long memoryRequiredInBytes = 0L; - try (TsFileSequenceReader sequenceReader = - new TsFileSequenceReader(tsFile.getPath(), true, true)) { - deviceMeasurementsMap = sequenceReader.getDeviceMeasurementsMap(); - memoryRequiredInBytes += - PipeMemoryWeightUtil.memoryOfIDeviceID2StrList(deviceMeasurementsMap); - - if (Objects.isNull(deviceIsAlignedMap)) { - deviceIsAlignedMap = new HashMap<>(); + try { + try (TsFileSequenceReader sequenceReader = + new TsFileSequenceReader(tsFile.getPath(), true, true)) { + cachedDeviceMeasurementsMap = sequenceReader.getDeviceMeasurementsMap(); + memoryRequiredInBytes += + PipeMemoryWeightUtil.memoryOfIDeviceID2StrList(cachedDeviceMeasurementsMap); + + cachedDeviceIsAlignedMap = new HashMap<>(); final TsFileDeviceIterator deviceIsAlignedIterator = sequenceReader.getAllDevicesIteratorWithIsAligned(); while (deviceIsAlignedIterator.hasNext()) { final Pair deviceIsAlignedPair = deviceIsAlignedIterator.next(); - deviceIsAlignedMap.put(deviceIsAlignedPair.getLeft(), deviceIsAlignedPair.getRight()); + cachedDeviceIsAlignedMap.put( + deviceIsAlignedPair.getLeft(), deviceIsAlignedPair.getRight()); } - } - memoryRequiredInBytes += PipeMemoryWeightUtil.memoryOfIDeviceId2Bool(deviceIsAlignedMap); + memoryRequiredInBytes += + PipeMemoryWeightUtil.memoryOfIDeviceId2Bool(cachedDeviceIsAlignedMap); - measurementDataTypeMap = sequenceReader.getFullPathDataTypeMap(); - memoryRequiredInBytes += PipeMemoryWeightUtil.memoryOfStr2TSDataType(measurementDataTypeMap); + cachedMeasurementDataTypeMap = sequenceReader.getFullPathDataTypeMap(); + memoryRequiredInBytes += + PipeMemoryWeightUtil.memoryOfStr2TSDataType(cachedMeasurementDataTypeMap); + } + } finally { + // The reader block is temporary and must be released even when metadata traversal fails. + readerMemoryBlock.close(); } - // Release memory of TsFileSequenceReader. - allocatedMemoryBlock.close(); - allocatedMemoryBlock = null; // Allocate again for the cached objects. - allocatedMemoryBlock = + final PipeMemoryBlock cachedMemoryBlock = PipeDataNodeResourceManager.memory() .forceAllocateIfSufficient(memoryRequiredInBytes, MEMORY_SUFFICIENT_THRESHOLD); - if (allocatedMemoryBlock == null) { + if (cachedMemoryBlock == null) { LOGGER.info( "PipeTsFileResource: Failed to cache objects for tsfile {} in cache, because memory usage is high", tsFile.getPath()); - deviceIsAlignedMap = null; - deviceMeasurementsMap = null; - measurementDataTypeMap = null; return false; } + // Publish all metadata only after the persistent accounting block is ready. + deviceMeasurementsMap = cachedDeviceMeasurementsMap; + deviceIsAlignedMap = cachedDeviceIsAlignedMap; + measurementDataTypeMap = cachedMeasurementDataTypeMap; + allocatedMemoryBlock = cachedMemoryBlock; LOGGER.info("PipeTsFileResource: Cached objects for tsfile {}.", tsFile.getPath()); return true; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFileResourceManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFileResourceManager.java index 2587e4a7cc4a..0dfefda76e49 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFileResourceManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFileResourceManager.java @@ -150,9 +150,9 @@ private File increaseFileReference( } try { increasePublicReference(resultFile, pipeName, isTsFile); - } catch (final IOException e) { + } catch (final IOException | RuntimeException e) { // The private reference must not outlive a failed public reference increase. - decreaseFileReference(resultFile, pipeName, false); + rollbackFileReference(resultFile, pipeName, e); throw e; } return resultFile; @@ -172,10 +172,27 @@ private boolean increaseReferenceIfExists( } finally { segmentLock.unlock(file); } - increasePublicReference(file, pipeName, isTsFile); + try { + increasePublicReference(file, pipeName, isTsFile); + } catch (final IOException | RuntimeException e) { + // The private reference is acquired before the public (assigner) reference. If the latter + // fails, roll back the reference acquired above; otherwise every failed retry permanently + // pins the pipe file in the logical memory/file pool. + rollbackFileReference(file, pipeName, e); + throw e; + } return true; } + private void rollbackFileReference( + final File file, final @Nullable String pipeName, final Exception originalException) { + try { + decreaseFileReference(file, pipeName, false); + } catch (final RuntimeException rollbackException) { + originalException.addSuppressed(rollbackException); + } + } + private void increasePublicReference( final File file, final String pipeName, final boolean isTsFile) throws IOException { if (Objects.isNull(pipeName)) { @@ -413,23 +430,62 @@ public Map getMeasurementDataTypeMapFromCache( public void pinTsFileResource( final TsFileResource resource, final boolean withMods, final @Nullable String pipeName) throws IOException { - increaseFileReference(resource.getTsFile(), true, pipeName); - if (withMods && resource.getModFile().exists()) { - // Avoid mod compaction - synchronized (resource.getModFile()) { - increaseFileReference(new File(resource.getModFile().getFilePath()), false, pipeName); + final File pinnedTsFile = increaseFileReference(resource.getTsFile(), true, pipeName); + try { + final ModificationFile modFile = resource.getModFile(); + if (withMods && modFile.exists()) { + // Avoid mod compaction. + synchronized (modFile) { + increaseFileReference(new File(modFile.getFilePath()), false, pipeName); + } + } + } catch (final IOException | RuntimeException e) { + // Pinning is a two-file operation. Do not leave the TsFile pinned when the mod file cannot + // be pinned (for example, when the pipe directory is temporarily unavailable). + try { + decreaseFileReference(pinnedTsFile, pipeName); + } catch (final RuntimeException rollbackException) { + e.addSuppressed(rollbackException); } + throw e; } } public void unpinTsFileResource(final TsFileResource resource, final @Nullable String pipeName) throws IOException { - final File pinnedFile = getHardlinkOrCopiedFileInPipeDir(resource.getTsFile(), pipeName); - decreaseFileReference(pinnedFile, pipeName); + Exception firstException = null; + File pinnedFile = null; + try { + pinnedFile = getHardlinkOrCopiedFileInPipeDir(resource.getTsFile(), pipeName); + decreaseFileReference(pinnedFile, pipeName); + } catch (final IOException | RuntimeException e) { + firstException = e; + } + + // Always attempt the mod-file cleanup even when resolving/decreasing the TsFile fails. A + // failed first cleanup must not strand the second reference. + try { + final File pinnedModFile = + pinnedFile != null + ? new File(pinnedFile + ModificationFile.FILE_SUFFIX) + : getHardlinkOrCopiedFileInPipeDir( + new File(resource.getModFile().getFilePath()), pipeName); + if (pinnedModFile.exists()) { + decreaseFileReference(pinnedModFile, pipeName); + } + } catch (final IOException | RuntimeException e) { + if (firstException == null) { + firstException = e; + } else { + firstException.addSuppressed(e); + } + } - final File modFile = new File(pinnedFile + ModificationFile.FILE_SUFFIX); - if (modFile.exists()) { - decreaseFileReference(modFile, pipeName); + if (firstException != null) { + if (firstException instanceof IOException) { + throw (IOException) firstException; + } + throw (RuntimeException) firstException; } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventBatch.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventBatch.java index 3e0ec9f779ef..de6055facfd2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventBatch.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventBatch.java @@ -19,10 +19,12 @@ package org.apache.iotdb.db.pipe.sink.payload.evolvable.batch; +import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; import org.apache.iotdb.commons.pipe.agent.task.progress.CommitterKey; import org.apache.iotdb.commons.pipe.event.EnrichedEvent; import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager; -import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryBlock; +import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryManager; +import org.apache.iotdb.db.pipe.resource.memory.PipeTabletMemoryBlock; import org.apache.iotdb.db.pipe.sink.protocol.thrift.async.IoTDBDataRegionAsyncSink; import org.apache.iotdb.db.storageengine.dataregion.wal.exception.WALPipeException; import org.apache.iotdb.pipe.api.event.Event; @@ -48,7 +50,8 @@ public abstract class PipeTabletEventBatch implements AutoCloseable { private long firstEventProcessingTime = Long.MIN_VALUE; protected long totalBufferSize = 0; - private final PipeMemoryBlock allocatedMemoryBlock; + private final PipeTabletMemoryBlock allocatedMemoryBlock; + private boolean shouldEmitOnMemoryPressure = false; protected volatile boolean isClosed = false; @@ -60,7 +63,8 @@ protected PipeTabletEventBatch( // limit in buffer size this.maxBatchSizeInBytes = requestMaxBatchSizeInBytes; - this.allocatedMemoryBlock = PipeDataNodeResourceManager.memory().forceAllocate(0); + this.allocatedMemoryBlock = + PipeDataNodeResourceManager.memory().forceAllocateForTabletWithRetry(0); if (recordMetric != null) { this.recordMetric = recordMetric; } else { @@ -90,15 +94,38 @@ public synchronized boolean onEvent(final TabletInsertionEvent event) if (((EnrichedEvent) event) .increaseReferenceCount(PipeTransferBatchReqBuilder.class.getName())) { + final int previousEventsSize = events.size(); + final long previousTotalBufferSize = totalBufferSize; + final boolean previousMemoryPressureState = shouldEmitOnMemoryPressure; + final Object batchState = captureBatchState(); + try { if (constructBatch(event)) { events.add((EnrichedEvent) event); } } catch (final Exception e) { - if (events.isEmpty()) { - clearBatchData(); - resetMemoryUsage(); + try { + rollbackBatchState(batchState); + } catch (final Exception rollbackException) { + e.addSuppressed(rollbackException); + } + + // A failed constructBatch must not retain a partial payload or its memory reservation, + // even when older events are already buffered in this batch. + if (totalBufferSize != previousTotalBufferSize) { + // Shrinking never needs to wait. This path still holds the batch monitor, so a + // blocking resize here would recreate the same lock cycle as a failed append. + PipeDataNodeResourceManager.memory() + .tryResize(allocatedMemoryBlock, previousTotalBufferSize); + } + totalBufferSize = previousTotalBufferSize; + shouldEmitOnMemoryPressure = + previousMemoryPressureState + || (e instanceof PipeRuntimeOutOfMemoryCriticalException && !events.isEmpty()); + if (events.size() > previousEventsSize) { + events.subList(previousEventsSize, events.size()).clear(); } + // If the event is not added to the batch, we need to decrease the reference count. ((EnrichedEvent) event) .decreaseReferenceCount(PipeTransferBatchReqBuilder.class.getName(), false); @@ -128,13 +155,34 @@ public synchronized boolean onEvent(final TabletInsertionEvent event) protected abstract boolean constructBatch(final TabletInsertionEvent event) throws WALPipeException, IOException; + /** Captures subclass payload state before constructing one event. */ + protected Object captureBatchState() { + return null; + } + + /** Restores subclass payload state after a failed event construction. */ + protected void rollbackBatchState(final Object state) {} + protected void increaseTotalBufferSizeAndUpdateMemoryBlock(final long bufferSize) { if (bufferSize <= 0) { return; } final long newTotalBufferSize = Math.min(totalBufferSize + bufferSize, maxBatchSizeInBytes); - PipeDataNodeResourceManager.memory().forceResize(allocatedMemoryBlock, newTotalBufferSize); + final PipeMemoryManager memoryManager = PipeDataNodeResourceManager.memory(); + if (!memoryManager.tryResize(allocatedMemoryBlock, newTotalBufferSize)) { + // Existing events must be emitted before this event is retried. Do not wait here: onEvent() + // is called while both the request builder and this batch are locked. + shouldEmitOnMemoryPressure = !events.isEmpty(); + throw new PipeRuntimeOutOfMemoryCriticalException( + String.format( + "forceResize: failed to allocate memory after %d retries, total memory size %d bytes, " + + "used memory size %d bytes, requested memory size %d bytes", + 0, + memoryManager.getTotalNonFloatingMemorySizeInBytes(), + memoryManager.getUsedMemorySizeInBytes(), + newTotalBufferSize - totalBufferSize)); + } totalBufferSize = newTotalBufferSize; } @@ -144,13 +192,20 @@ protected void releaseAllocatedMemoryBlock() { protected void clearBatchData() {} + /** Close resources owned by a batch that will not be reused. */ + protected void closeBatchData() { + clearBatchData(); + } + public boolean shouldEmit() { if (events.isEmpty()) { return false; } final long diff = System.currentTimeMillis() - firstEventProcessingTime; - if (totalBufferSize >= maxBatchSizeInBytes || diff >= maxDelayInMs) { + if (shouldEmitOnMemoryPressure + || totalBufferSize >= maxBatchSizeInBytes + || diff >= maxDelayInMs) { recordMetric.accept(diff, totalBufferSize, events.size()); return true; } @@ -159,8 +214,20 @@ public boolean shouldEmit() { public synchronized void onSuccess() { events.clear(); + try { + clearBatchData(); + } finally { + resetMemoryUsage(); + } + } - resetMemoryUsage(); + /** + * Close a detached asynchronous batch after its event references have been handed to handlers or + * a retry queue. Unlike {@link #close()}, this method does not release those references. + */ + public synchronized void closeAfterEventTransfer() { + events.clear(); + close(); } @Override @@ -170,11 +237,20 @@ public synchronized void close() { } isClosed = true; - clearEventsReferenceCount(PipeTabletEventBatch.class.getName()); - events.clear(); - clearBatchData(); - resetMemoryUsage(); - allocatedMemoryBlock.close(); + try { + clearEventsReferenceCount(PipeTabletEventBatch.class.getName()); + } finally { + events.clear(); + try { + closeBatchData(); + } finally { + try { + resetMemoryUsage(); + } finally { + allocatedMemoryBlock.close(); + } + } + } } /** @@ -204,6 +280,7 @@ public synchronized void discardEventsOfPipe(final CommitterKey committerKey) { private void resetMemoryUsage() { totalBufferSize = 0; + shouldEmitOnMemoryPressure = false; releaseAllocatedMemoryBlock(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventPlainBatch.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventPlainBatch.java index c802526bf47a..a99564708fc3 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventPlainBatch.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventPlainBatch.java @@ -64,18 +64,51 @@ protected boolean constructBatch(final TabletInsertionEvent event) throws IOExce } @Override - public synchronized void onSuccess() { - clearBatchData(); + protected void clearBatchData() { + insertNodeBuffers.clear(); + tabletBuffers.clear(); - super.onSuccess(); + pipe2BytesAccumulated.clear(); } @Override - protected void clearBatchData() { - insertNodeBuffers.clear(); - tabletBuffers.clear(); + protected Object captureBatchState() { + return new BatchState( + insertNodeBuffers.size(), tabletBuffers.size(), new HashMap<>(pipe2BytesAccumulated)); + } + + @Override + protected void rollbackBatchState(final Object state) { + if (!(state instanceof BatchState)) { + return; + } + final BatchState batchState = (BatchState) state; + truncate(insertNodeBuffers, batchState.insertNodeBuffersSize); + truncate(tabletBuffers, batchState.tabletBuffersSize); pipe2BytesAccumulated.clear(); + pipe2BytesAccumulated.putAll(batchState.pipe2BytesAccumulated); + } + + private static void truncate(final List list, final int size) { + if (list.size() > size) { + list.subList(size, list.size()).clear(); + } + } + + private static final class BatchState { + private final int insertNodeBuffersSize; + private final int tabletBuffersSize; + private final Map, Long> pipe2BytesAccumulated; + + private BatchState( + final int insertNodeBuffersSize, + final int tabletBuffersSize, + final Map, Long> pipe2BytesAccumulated) { + this.insertNodeBuffersSize = insertNodeBuffersSize; + this.tabletBuffersSize = tabletBuffersSize; + this.pipe2BytesAccumulated = pipe2BytesAccumulated; + } } public PipeTransferTabletBatchReq toTPipeTransferReq() throws IOException { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventTsFileBatch.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventTsFileBatch.java index 6aa6518a652f..1f04fe8a904d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventTsFileBatch.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventTsFileBatch.java @@ -156,24 +156,36 @@ protected boolean constructBatch(final TabletInsertionEvent event) { final PipeInsertNodeTabletInsertionEvent insertNodeTabletInsertionEvent = (PipeInsertNodeTabletInsertionEvent) event; final List tablets = insertNodeTabletInsertionEvent.convertToTablets(); - increaseTotalBufferSizeAndUpdateMemoryBlock(calculateTabletsSizeInBytes(tablets)); + final List retainedTablets = new ArrayList<>(tablets.size()); + final List retainedAlignedFlags = new ArrayList<>(tablets.size()); for (int i = 0; i < tablets.size(); ++i) { final Tablet tablet = tablets.get(i); if (tablet.rowSize == 0) { continue; } + retainedTablets.add(tablet); + retainedAlignedFlags.add(insertNodeTabletInsertionEvent.isAligned(i)); + } + + if (retainedTablets.isEmpty()) { + return false; + } + increaseTotalBufferSizeAndUpdateMemoryBlock(calculateTabletsSizeInBytes(retainedTablets)); + for (int i = 0; i < retainedTablets.size(); ++i) { + final Tablet tablet = retainedTablets.get(i); bufferTablet( insertNodeTabletInsertionEvent.getPipeName(), insertNodeTabletInsertionEvent.getCreationTime(), tablet, - insertNodeTabletInsertionEvent.isAligned(i)); + retainedAlignedFlags.get(i)); } + return true; } else if (event instanceof PipeRawTabletInsertionEvent) { final PipeRawTabletInsertionEvent rawTabletInsertionEvent = (PipeRawTabletInsertionEvent) event; final Tablet tablet = rawTabletInsertionEvent.convertToTablet(); if (tablet.rowSize == 0) { - return true; + return false; } increaseTotalBufferSizeAndUpdateMemoryBlock(calculateTabletSizeInBytes(tablet)); bufferTablet( @@ -181,6 +193,7 @@ protected boolean constructBatch(final TabletInsertionEvent event) { rawTabletInsertionEvent.getCreationTime(), tablet, rawTabletInsertionEvent.isAligned()); + return true; } else { LOGGER.warn( "Batch id = {}: Unsupported event {} type {} when constructing tsfile batch", @@ -188,7 +201,7 @@ protected boolean constructBatch(final TabletInsertionEvent event) { event, event.getClass()); } - return true; + return false; } private long calculateTabletsSizeInBytes(final List tablets) { @@ -202,6 +215,45 @@ private static long calculateTabletSizeInBytes(final Tablet tablet) { return PipeMemoryWeightUtil.calculateTabletSizeInBytes(tablet) * 2; } + @Override + public Object captureBatchState() { + return new BatchState( + tabletList.size(), isTabletAlignedList.size(), new HashMap<>(pipeName2WeightMap)); + } + + @Override + public void rollbackBatchState(final Object state) { + if (!(state instanceof BatchState)) { + return; + } + final BatchState batchState = (BatchState) state; + truncate(tabletList, batchState.tabletListSize); + truncate(isTabletAlignedList, batchState.isTabletAlignedListSize); + pipeName2WeightMap.clear(); + pipeName2WeightMap.putAll(batchState.pipeName2WeightMap); + } + + private static void truncate(final List list, final int size) { + if (list.size() > size) { + list.subList(size, list.size()).clear(); + } + } + + private static final class BatchState { + private final int tabletListSize; + private final int isTabletAlignedListSize; + private final Map, Double> pipeName2WeightMap; + + private BatchState( + final int tabletListSize, + final int isTabletAlignedListSize, + final Map, Double> pipeName2WeightMap) { + this.tabletListSize = tabletListSize; + this.isTabletAlignedListSize = isTabletAlignedListSize; + this.pipeName2WeightMap = pipeName2WeightMap; + } + } + private void bufferTablet( final String pipeName, final long creationTime, @@ -286,7 +338,13 @@ private List writeTabletsToTsFiles() throws IOException, WriteProcessExcep // Try making the tsfile size as large as possible while (!device2TabletsLinkedList.isEmpty()) { if (Objects.isNull(fileWriter)) { - fileWriter = new TsFileWriter(createFile()); + final File file = createFile(); + try { + fileWriter = new TsFileWriter(file); + } catch (final IOException | RuntimeException e) { + FileUtils.deleteQuietly(file); + throw e; + } } try { @@ -472,13 +530,6 @@ private void tryBestToWriteTabletsIntoOneFile( } } - @Override - public synchronized void onSuccess() { - clearBatchData(); - - super.onSuccess(); - } - @Override protected void clearBatchData() { pipeName2WeightMap.clear(); @@ -492,34 +543,37 @@ protected void clearBatchData() { } @Override - public synchronized void close() { + protected void closeBatchData() { + pipeName2WeightMap.clear(); + tabletList.clear(); + isTabletAlignedList.clear(); + if (Objects.nonNull(fileWriter)) { + final File file = fileWriter.getIOWriter().getFile(); try { fileWriter.close(); } catch (final Exception e) { LOGGER.info( "Batch id = {}: Failed to close the tsfile {} when trying to close batch, because {}", currentBatchId.get(), - fileWriter.getIOWriter().getFile().getPath(), + file.getPath(), e.getMessage(), e); } try { - RetryUtils.retryOnException(() -> FileUtils.delete(fileWriter.getIOWriter().getFile())); + RetryUtils.retryOnException(() -> FileUtils.delete(file)); } catch (final Exception e) { LOGGER.info( "Batch id = {}: Failed to delete the tsfile {} when trying to close batch, because {}", currentBatchId.get(), - fileWriter.getIOWriter().getFile().getPath(), + file.getPath(), e.getMessage(), e); + } finally { + fileWriter = null; } - - fileWriter = null; } - - super.close(); } protected File createFile() throws IOException { @@ -546,10 +600,17 @@ private List writeTabletsToTsFiles() throws org.apache.iotdb.db.exception.WriteProcessException { final IMemTable memTable = new PrimitiveMemTable(null, null); final List sealedFiles = new ArrayList<>(); - try (final RestorableTsFileIOWriter writer = new RestorableTsFileIOWriter(createFile())) { - writeTabletsIntoOneFile(memTable, writer); - sealedFiles.add(writer.getFile()); + File file = null; + try { + file = createFile(); + try (final RestorableTsFileIOWriter writer = new RestorableTsFileIOWriter(file)) { + writeTabletsIntoOneFile(memTable, writer); + sealedFiles.add(writer.getFile()); + } } catch (final Exception e) { + if (file != null) { + FileUtils.deleteQuietly(file); + } LOGGER.warn( "Batch id = {}: Failed to write tablets into tsfile, because {}", currentBatchId.get(), diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilder.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilder.java index 52827ac00088..3921d42e4222 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilder.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilder.java @@ -66,6 +66,7 @@ public class PipeTransferBatchReqBuilder implements AutoCloseable { private static final Logger LOGGER = LoggerFactory.getLogger(PipeTransferBatchReqBuilder.class); + private final boolean usingTsFileBatch; private final boolean useLeaderCache; private final int requestMaxDelayInMs; @@ -80,14 +81,14 @@ public class PipeTransferBatchReqBuilder implements AutoCloseable { // If the leader cache is disabled (or unable to find the endpoint of event in the leader cache), // the event will be stored in the default batch. - private final PipeTabletEventBatch defaultBatch; + private PipeTabletEventBatch defaultBatch; // If the leader cache is enabled, the batch will be divided by the leader endpoint, // each endpoint has a batch. // This is only used in plain batch since tsfile does not return redirection info. private final Map endPointToBatch = new HashMap<>(); public PipeTransferBatchReqBuilder(final PipeParameters parameters) { - final boolean usingTsFileBatch = + usingTsFileBatch = parameters .getStringOrDefault( Arrays.asList(CONNECTOR_FORMAT_KEY, SINK_FORMAT_KEY), CONNECTOR_FORMAT_HYBRID_VALUE) @@ -119,12 +120,20 @@ public PipeTransferBatchReqBuilder(final PipeParameters parameters) { usingTsFileBatch ? CONNECTOR_IOTDB_TS_FILE_BATCH_SIZE_DEFAULT_VALUE : CONNECTOR_IOTDB_PLAIN_BATCH_SIZE_DEFAULT_VALUE); - this.defaultBatch = - usingTsFileBatch - ? new PipeTabletEventTsFileBatch( - requestMaxDelayInMs, requestMaxBatchSizeInBytes, this::recordTsFileMetric) - : new PipeTabletEventPlainBatch( - requestMaxDelayInMs, requestMaxBatchSizeInBytes, this::recordTabletMetric); + this.defaultBatch = createDefaultBatch(); + } + + private PipeTabletEventBatch createDefaultBatch() { + return usingTsFileBatch + ? new PipeTabletEventTsFileBatch( + requestMaxDelayInMs, requestMaxBatchSizeInBytes, this::recordTsFileMetric) + : new PipeTabletEventPlainBatch( + requestMaxDelayInMs, requestMaxBatchSizeInBytes, this::recordTabletMetric); + } + + private PipeTabletEventPlainBatch createLeaderCacheBatch() { + return new PipeTabletEventPlainBatch( + requestMaxDelayInMs, requestMaxBatchSizeInBytes, this::recordTabletMetric); } /** @@ -190,6 +199,59 @@ public synchronized void onEvent(final TabletInsertionEvent event) return nonEmptyAndShouldEmitBatches; } + /** + * Atomically detaches every batch that is ready to emit. + * + *

The detached batches are immutable from the builder's point of view: subsequent events are + * appended to fresh batches. This is required for asynchronous sinks, whose completion callback + * may clear a batch after another sink thread has already appended a new event to it. + */ + public synchronized List> + getAllNonEmptyAndShouldEmitBatchesAndDetach() { + final List> batches = + new ArrayList<>(endPointToBatch.size() + 1); + if (!defaultBatch.isEmpty() && defaultBatch.shouldEmit()) { + batches.add(new Pair<>(null, defaultBatch)); + } + + for (final Map.Entry entry : endPointToBatch.entrySet()) { + final PipeTabletEventPlainBatch batch = entry.getValue(); + if (!batch.isEmpty() && batch.shouldEmit()) { + batches.add(new Pair<>(entry.getKey(), batch)); + } + } + + // Construct all replacement batches before changing the builder mappings. If construction of + // one replacement fails (for example while creating a TsFile batch directory), the old + // batches remain owned by this builder and can still be retried or closed by the caller. + final List replacements = new ArrayList<>(batches.size()); + try { + for (final Pair batch : batches) { + replacements.add(batch.getLeft() == null ? createDefaultBatch() : createLeaderCacheBatch()); + } + } catch (final RuntimeException | Error e) { + replacements.forEach( + replacement -> { + try { + replacement.close(); + } catch (final RuntimeException | Error closeException) { + e.addSuppressed(closeException); + } + }); + throw e; + } + + for (int i = 0; i < batches.size(); ++i) { + final Pair batch = batches.get(i); + if (batch.getLeft() == null) { + defaultBatch = replacements.get(i); + } else { + endPointToBatch.put(batch.getLeft(), (PipeTabletEventPlainBatch) replacements.get(i)); + } + } + return batches; + } + public synchronized boolean isEmpty() { if (!defaultBatch.isEmpty()) { return false; @@ -202,6 +264,23 @@ public synchronized boolean isEmpty() { return true; } + /** Returns whether a specific event is still retained by one of the current batches. */ + public synchronized boolean containsEvent(final Event event) { + for (final EnrichedEvent batchedEvent : defaultBatch.events) { + if (batchedEvent == event) { + return true; + } + } + for (final PipeTabletEventPlainBatch batch : endPointToBatch.values()) { + for (final EnrichedEvent batchedEvent : batch.events) { + if (batchedEvent == event) { + return true; + } + } + } + return false; + } + public synchronized void discardEventsOfPipe( final String pipeNameToDrop, final long creationTimeToDrop, final int regionId) { discardEventsOfPipe(new CommitterKey(pipeNameToDrop, creationTimeToDrop, regionId, -1)); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSink.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSink.java index ade003240847..648c3f4e7403 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSink.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSink.java @@ -20,6 +20,7 @@ package org.apache.iotdb.db.pipe.sink.protocol.airgap; import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; import org.apache.iotdb.commons.pipe.agent.task.progress.CommitterKey; import org.apache.iotdb.commons.pipe.event.EnrichedEvent; import org.apache.iotdb.commons.pipe.sink.limiter.TsFileSendRateLimiter; @@ -125,7 +126,17 @@ public void transfer(final TabletInsertionEvent tabletInsertionEvent) throws Exc // We need to restore the transfer quickly by retry under this circumstance socket.setSoTimeout(PIPE_CONFIG.getPipeAirGapSinkTabletTimeoutMs()); if (isTabletBatchModeEnabled) { - tabletBatchBuilder.onEvent(tabletInsertionEvent); + try { + tabletBatchBuilder.onEvent(tabletInsertionEvent); + } catch (final PipeRuntimeOutOfMemoryCriticalException memoryException) { + try { + doTransferWrapper(socket); + } catch (final Exception transferException) { + transferException.addSuppressed(memoryException); + throw transferException; + } + tabletBatchBuilder.onEvent(tabletInsertionEvent); + } doTransferWrapper(socket); } else if (tabletInsertionEvent instanceof PipeInsertNodeTabletInsertionEvent) { doTransferWrapper(socket, (PipeInsertNodeTabletInsertionEvent) tabletInsertionEvent); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/pipeconsensus/payload/builder/PipeConsensusTransferBatchReqBuilder.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/pipeconsensus/payload/builder/PipeConsensusTransferBatchReqBuilder.java index 46968f64a798..85b6918abe14 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/pipeconsensus/payload/builder/PipeConsensusTransferBatchReqBuilder.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/pipeconsensus/payload/builder/PipeConsensusTransferBatchReqBuilder.java @@ -21,6 +21,7 @@ import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId; import org.apache.iotdb.commons.consensus.index.ProgressIndex; +import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; import org.apache.iotdb.commons.pipe.event.EnrichedEvent; import org.apache.iotdb.consensus.pipe.thrift.TCommitId; import org.apache.iotdb.consensus.pipe.thrift.TPipeConsensusTransferReq; @@ -162,7 +163,16 @@ private void increaseTotalBufferSizeAndUpdateMemoryBlock(final long bufferSize) final long newTotalBufferSize = Math.min(totalBufferSize + bufferSize, getMaxBatchSizeInBytes()); - PipeDataNodeResourceManager.memory().forceResize(allocatedMemoryBlock, newTotalBufferSize); + if (!PipeDataNodeResourceManager.memory().tryResize(allocatedMemoryBlock, newTotalBufferSize)) { + throw new PipeRuntimeOutOfMemoryCriticalException( + String.format( + "forceResize: failed to allocate memory after %d retries, total memory size %d bytes, " + + "used memory size %d bytes, requested memory size %d bytes", + 0, + PipeDataNodeResourceManager.memory().getTotalNonFloatingMemorySizeInBytes(), + PipeDataNodeResourceManager.memory().getUsedMemorySizeInBytes(), + newTotalBufferSize - totalBufferSize)); + } totalBufferSize = newTotalBufferSize; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java index 1387b5fbdb75..e6221f6706b6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java @@ -23,6 +23,7 @@ import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.client.ThriftClient; import org.apache.iotdb.commons.client.async.AsyncPipeDataTransferServiceClient; +import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; import org.apache.iotdb.commons.exception.pipe.PipeRuntimeSinkNonReportTimeConfigurableException; import org.apache.iotdb.commons.exception.pipe.PipeRuntimeSinkResourceException; import org.apache.iotdb.commons.pipe.agent.task.progress.CommitterKey; @@ -124,11 +125,16 @@ public class IoTDBDataRegionAsyncSink extends IoTDBSink { private final BlockingQueue retryTsFileQueue = new LinkedBlockingQueue<>(); private final PipeDataRegionEventCounter retryEventQueueEventCounter = new PipeDataRegionEventCounter(); - // Guarded by this. Events need identity semantics because the same payload may compare equal. + // Guarded by this. The map is also the retry-queue membership index. Events need identity + // semantics because the same payload may compare equal. private final Map retryEvent2ResourceFailureType = new IdentityHashMap<>(); // Keep only the latest text to avoid retaining the complete exception chain for every event. private volatile String lastRetryFailureMessage; + // Events removed from the retry queue remain here while their next transfer is being started. + // A callback from an older handler must not create a second queue entry. + private final Map> retryingEvent2Handlers = + new IdentityHashMap<>(); private IoTDBDataNodeAsyncClientManager clientManager; private IoTDBDataNodeAsyncClientManager transferTsFileClientManager; @@ -238,13 +244,31 @@ public void transfer(final TabletInsertionEvent tabletInsertionEvent) throws Exc } if (isTabletBatchModeEnabled) { - tabletBatchBuilder.onEvent(tabletInsertionEvent); - transferBatchedEventsIfNecessary(); + addTabletEventToBatchAndTransferIfNecessary(tabletInsertionEvent); } else { transferInEventWithoutCheck(tabletInsertionEvent); } } + private void addTabletEventToBatchAndTransferIfNecessary( + final TabletInsertionEvent tabletInsertionEvent) throws Exception { + try { + tabletBatchBuilder.onEvent(tabletInsertionEvent); + } catch (final PipeRuntimeOutOfMemoryCriticalException memoryException) { + // The failed append was rolled back. Flush buffered events, then retry the current event + // after + // the detached batch has released its memory. + try { + transferBatchedEventsIfNecessary(); + } catch (final Exception transferException) { + transferException.addSuppressed(memoryException); + throw transferException; + } + tabletBatchBuilder.onEvent(tabletInsertionEvent); + } + transferBatchedEventsIfNecessary(); + } + private void transferInBatchWithoutCheck( final Pair endPointAndBatch) throws IOException, WriteProcessException { @@ -254,50 +278,100 @@ private void transferInBatchWithoutCheck( final PipeTabletEventBatch batch = endPointAndBatch.getRight(); if (batch instanceof PipeTabletEventPlainBatch) { - transfer( - endPointAndBatch.getLeft(), - new PipeTransferTabletBatchEventHandler((PipeTabletEventPlainBatch) batch, this)); + transferInPlainBatchWithoutCheck( + endPointAndBatch.getLeft(), (PipeTabletEventPlainBatch) batch); } else if (batch instanceof PipeTabletEventTsFileBatch) { - final PipeTabletEventTsFileBatch tsFileBatch = (PipeTabletEventTsFileBatch) batch; - final List sealedFiles = tsFileBatch.sealTsFiles(); - final Map, Double> pipe2WeightMap = tsFileBatch.deepCopyPipe2WeightMap(); - final List events = tsFileBatch.deepCopyEvents(); - final AtomicInteger eventsReferenceCount = new AtomicInteger(sealedFiles.size()); - final AtomicBoolean eventsHadBeenAddedToRetryQueue = new AtomicBoolean(false); + transferInTsFileBatchWithoutCheck((PipeTabletEventTsFileBatch) batch); + } else { + final Exception exception = + new PipeException( + String.format( + "Unsupported batch type %s when transferring tablet insertion event.", + batch.getClass())); + addFailureEventsToRetryQueue(batch.deepCopyEvents(), exception); + batch.closeAfterEventTransfer(); + } + } - int transferredFileCount = 0; - try { - for (int outputIndex = 0; outputIndex < sealedFiles.size(); outputIndex++) { - final File sealedFile = sealedFiles.get(outputIndex); - transfer( - new PipeTransferTsFileHandler( - this, - pipe2WeightMap, - events, - eventsReferenceCount, - eventsHadBeenAddedToRetryQueue, - sealedFile, - null, - false, - null, - outputIndex)); - transferredFileCount++; - } - } catch (final Exception e) { - for (int i = transferredFileCount; i < sealedFiles.size(); i++) { - FileUtils.deleteQuietly(sealedFiles.get(i)); - } - PipeLogger.log(LOGGER::warn, e, "Failed to transfer tsfile batch (%s).", sealedFiles); - if (eventsHadBeenAddedToRetryQueue.compareAndSet(false, true)) { - addFailureEventsToRetryQueue(events, e); - } + private void transferInPlainBatchWithoutCheck( + final TEndPoint endPoint, final PipeTabletEventPlainBatch batch) { + final List events = batch.deepCopyEvents(); + boolean isBatchClosed = false; + try { + final PipeTransferTabletBatchEventHandler handler = + new PipeTransferTabletBatchEventHandler(batch, this); + trackRetryHandler(handler, events); + // The handler owns a request snapshot and the event references now. Release batch memory + // before borrowing a client. + isBatchClosed = true; + batch.closeAfterEventTransfer(); + transfer(endPoint, handler); + } catch (final Exception e) { + addFailureEventsToRetryQueue(events, e); + PipeLogger.log( + LOGGER::warn, + e, + "Failed to transfer TabletInsertionEvent batch. Total failed events: %s.", + events.size()); + } finally { + if (!isBatchClosed) { + batch.closeAfterEventTransfer(); } - } else { - LOGGER.warn( - "Unsupported batch type {} when transferring tablet insertion event.", batch.getClass()); } + } - endPointAndBatch.getRight().onSuccess(); + private void transferInTsFileBatchWithoutCheck(final PipeTabletEventTsFileBatch batch) { + final List events = batch.deepCopyEvents(); + final AtomicBoolean eventsHadBeenAddedToRetryQueue = new AtomicBoolean(false); + List sealedFiles = Collections.emptyList(); + int transferredFileCount = 0; + boolean isBatchClosed = false; + try { + sealedFiles = batch.sealTsFiles(); + if (sealedFiles.isEmpty()) { + throw new PipeException( + String.format( + "Failed to transfer tsfile batch because no tsfile was generated for %s.", batch)); + } + final Map, Double> pipe2WeightMap = batch.deepCopyPipe2WeightMap(); + final AtomicInteger eventsReferenceCount = new AtomicInteger(sealedFiles.size()); + + // Conversion produced self-contained files, so the detached batch can release tablet memory + // before handlers reserve their read buffers. + isBatchClosed = true; + batch.closeAfterEventTransfer(); + + for (int outputIndex = 0; outputIndex < sealedFiles.size(); outputIndex++) { + final File sealedFile = sealedFiles.get(outputIndex); + final PipeTransferTsFileHandler handler = + new PipeTransferTsFileHandler( + this, + pipe2WeightMap, + events, + eventsReferenceCount, + eventsHadBeenAddedToRetryQueue, + sealedFile, + null, + false, + null, + outputIndex); + trackRetryHandler(handler, events); + transfer(handler); + transferredFileCount++; + } + } catch (final Exception e) { + for (int i = transferredFileCount; i < sealedFiles.size(); i++) { + FileUtils.deleteQuietly(sealedFiles.get(i)); + } + PipeLogger.log(LOGGER::warn, e, "Failed to transfer tsfile batch (%s).", sealedFiles); + if (eventsHadBeenAddedToRetryQueue.compareAndSet(false, true)) { + addFailureEventsToRetryQueue(events, e); + } + } finally { + if (!isBatchClosed) { + batch.closeAfterEventTransfer(); + } + } } private boolean transferInEventWithoutCheck(final TabletInsertionEvent tabletInsertionEvent) @@ -310,17 +384,28 @@ private boolean transferInEventWithoutCheck(final TabletInsertionEvent tabletIns IoTDBDataRegionAsyncSink.class.getName())) { return false; } - - final InsertNode insertNode = pipeInsertNodeTabletInsertionEvent.getInsertNode(); - final TPipeTransferReq pipeTransferReq = - compressIfNeeded(PipeTransferTabletInsertNodeReq.toTPipeTransferReq(insertNode)); - final PipeTransferTabletInsertNodeEventHandler pipeTransferInsertNodeReqHandler = - new PipeTransferTabletInsertNodeEventHandler( - pipeInsertNodeTabletInsertionEvent, pipeTransferReq, this); - - transfer( - // getDeviceId() may return null for InsertRowsNode - pipeInsertNodeTabletInsertionEvent.getDeviceId(), pipeTransferInsertNodeReqHandler); + boolean handedToHandler = false; + try { + final InsertNode insertNode = pipeInsertNodeTabletInsertionEvent.getInsertNode(); + final TPipeTransferReq pipeTransferReq = + compressIfNeeded(PipeTransferTabletInsertNodeReq.toTPipeTransferReq(insertNode)); + final PipeTransferTabletInsertNodeEventHandler pipeTransferInsertNodeReqHandler = + new PipeTransferTabletInsertNodeEventHandler( + pipeInsertNodeTabletInsertionEvent, pipeTransferReq, this); + trackRetryHandler( + pipeTransferInsertNodeReqHandler, + Collections.singletonList(pipeInsertNodeTabletInsertionEvent)); + handedToHandler = true; + + transfer( + // getDeviceId() may return null for InsertRowsNode + pipeInsertNodeTabletInsertionEvent.getDeviceId(), pipeTransferInsertNodeReqHandler); + } finally { + if (!handedToHandler) { + pipeInsertNodeTabletInsertionEvent.decreaseReferenceCount( + IoTDBDataRegionAsyncSink.class.getName(), false); + } + } } else { // tabletInsertionEvent instanceof PipeRawTabletInsertionEvent final PipeRawTabletInsertionEvent pipeRawTabletInsertionEvent = (PipeRawTabletInsertionEvent) tabletInsertionEvent; @@ -329,17 +414,27 @@ private boolean transferInEventWithoutCheck(final TabletInsertionEvent tabletIns IoTDBDataRegionAsyncSink.class.getName())) { return false; } - - final TPipeTransferReq pipeTransferTabletRawReq = - compressIfNeeded( - PipeTransferTabletRawReq.toTPipeTransferReq( - pipeRawTabletInsertionEvent.convertToTablet(), - pipeRawTabletInsertionEvent.isAligned())); - final PipeTransferTabletRawEventHandler pipeTransferTabletReqHandler = - new PipeTransferTabletRawEventHandler( - pipeRawTabletInsertionEvent, pipeTransferTabletRawReq, this); - - transfer(pipeRawTabletInsertionEvent.getDeviceId(), pipeTransferTabletReqHandler); + boolean handedToHandler = false; + try { + final TPipeTransferReq pipeTransferTabletRawReq = + compressIfNeeded( + PipeTransferTabletRawReq.toTPipeTransferReq( + pipeRawTabletInsertionEvent.convertToTablet(), + pipeRawTabletInsertionEvent.isAligned())); + final PipeTransferTabletRawEventHandler pipeTransferTabletReqHandler = + new PipeTransferTabletRawEventHandler( + pipeRawTabletInsertionEvent, pipeTransferTabletRawReq, this); + trackRetryHandler( + pipeTransferTabletReqHandler, Collections.singletonList(pipeRawTabletInsertionEvent)); + handedToHandler = true; + + transfer(pipeRawTabletInsertionEvent.getDeviceId(), pipeTransferTabletReqHandler); + } finally { + if (!handedToHandler) { + pipeRawTabletInsertionEvent.decreaseReferenceCount( + IoTDBDataRegionAsyncSink.class.getName(), false); + } + } } return true; @@ -408,14 +503,14 @@ private boolean transferWithoutCheck(final TsFileInsertionEvent tsFileInsertionE return false; } - // We assume that no exceptions will be thrown after reference count is increased. + PipeTransferTsFileHandler pipeTransferTsFileHandler = null; try { // Just in case. To avoid the case that exception occurred when constructing the handler. if (!pipeTsFileInsertionEvent.getTsFile().exists()) { throw new FileNotFoundException(pipeTsFileInsertionEvent.getTsFile().getAbsolutePath()); } - final PipeTransferTsFileHandler pipeTransferTsFileHandler = + pipeTransferTsFileHandler = new PipeTransferTsFileHandler( this, Collections.singletonMap( @@ -431,10 +526,19 @@ private boolean transferWithoutCheck(final TsFileInsertionEvent tsFileInsertionE pipeTsFileInsertionEvent.isWithMod() && clientManager.supportModsIfIsDataNodeReceiver(), pipeTsFileInsertionEvent.getDatabaseName()); + trackRetryHandler( + pipeTransferTsFileHandler, Collections.singletonList(pipeTsFileInsertionEvent)); transfer(pipeTransferTsFileHandler); return true; } catch (final Exception e) { + if (pipeTransferTsFileHandler != null) { + try { + pipeTransferTsFileHandler.close(); + } catch (final RuntimeException closeException) { + e.addSuppressed(closeException); + } + } // Just in case. To avoid the case that exception occurred when constructing the handler. pipeTsFileInsertionEvent.decreaseReferenceCount( IoTDBDataRegionAsyncSink.class.getName(), false); @@ -464,7 +568,9 @@ private void transfer(final PipeTransferTsFileHandler pipeTransferTsFileHandler) transferTsFileClientManager.getExecutor()); } catch (final RuntimeException e) { transferTsFileCounter.decrementAndGet(); - throw e; + logOnClientException(null, e); + pipeTransferTsFileHandler.onError(e); + return; } if (PipeConfig.getInstance().isTransferTsFileSync()) { @@ -513,7 +619,7 @@ private void transferBatchedEventsIfNecessary() throws IOException, WriteProcess } for (final Pair endPointAndBatch : - tabletBatchBuilder.getAllNonEmptyAndShouldEmitBatches()) { + tabletBatchBuilder.getAllNonEmptyAndShouldEmitBatchesAndDetach()) { transferInBatchWithoutCheck(endPointAndBatch); } } @@ -571,6 +677,7 @@ private void transferQueuedEventsIfNecessary(final boolean forced) { final long retryStartTime = System.currentTimeMillis(); final int remainingEvents = retryEventQueue.size() + retryTsFileQueue.size(); while (!retryEventQueue.isEmpty() || !retryTsFileQueue.isEmpty()) { + final Event retryEvent; synchronized (this) { if (isClosed.get()) { return; @@ -579,44 +686,43 @@ private void transferQueuedEventsIfNecessary(final boolean forced) { break; } - final Event peekedEvent; - final Event polledEvent; if (!retryEventQueue.isEmpty()) { - peekedEvent = retryEventQueue.peek(); - retryEvent2ResourceFailureType.remove(peekedEvent); - - if (peekedEvent instanceof PipeInsertNodeTabletInsertionEvent) { - retryTransfer((PipeInsertNodeTabletInsertionEvent) peekedEvent); - } else if (peekedEvent instanceof PipeRawTabletInsertionEvent) { - retryTransfer((PipeRawTabletInsertionEvent) peekedEvent); - } else { - LOGGER.warn( - "IoTDBThriftAsyncConnector does not support transfer generic event: {}.", - peekedEvent); - } - - polledEvent = retryEventQueue.poll(); + retryEvent = retryEventQueue.poll(); } else { if (transferTsFileCounter.get() >= PipeConfig.getInstance().getPipeRealTimeQueueMaxWaitingTsFileSize()) { return; } - peekedEvent = retryTsFileQueue.peek(); - retryEvent2ResourceFailureType.remove(peekedEvent); - retryTransfer((PipeTsFileInsertionEvent) peekedEvent); - polledEvent = retryTsFileQueue.poll(); + retryEvent = retryTsFileQueue.poll(); } - retryEventQueueEventCounter.decreaseEventCount(polledEvent); - if (polledEvent != peekedEvent) { - LOGGER.error( - "The event polled from the queue is not the same as the event peeked from the queue. " - + "Peeked event: {}, polled event: {}.", - peekedEvent, - polledEvent); + if (retryEvent == null) { + break; + } + // Remove queue membership before retrying. The original queue reference remains owned by + // this in-flight attempt until the retry is handed to a batch or handler. + retryEvent2ResourceFailureType.remove(retryEvent); + retryEventQueueEventCounter.decreaseEventCount(retryEvent); + retryingEvent2Handlers.put(retryEvent, Collections.newSetFromMap(new IdentityHashMap<>())); + } + + // Handler callbacks call back into this sink, so transfer startup must happen outside the + // sink monitor. + if (retryEvent instanceof PipeInsertNodeTabletInsertionEvent) { + retryTransfer((PipeInsertNodeTabletInsertionEvent) retryEvent); + } else if (retryEvent instanceof PipeRawTabletInsertionEvent) { + retryTransfer((PipeRawTabletInsertionEvent) retryEvent); + } else if (retryEvent instanceof PipeTsFileInsertionEvent) { + retryTransfer((PipeTsFileInsertionEvent) retryEvent); + } else { + LOGGER.warn( + "IoTDBThriftAsyncConnector does not support transfer generic event: {}.", retryEvent); + synchronized (this) { + retryingEvent2Handlers.remove(retryEvent); } - if (polledEvent != null && LOGGER.isDebugEnabled()) { - LOGGER.debug("Polled event {} from retry queue.", polledEvent); + if (retryEvent instanceof EnrichedEvent) { + ((EnrichedEvent) retryEvent) + .clearReferenceCount(IoTDBDataRegionAsyncSink.class.getName()); } } @@ -663,12 +769,22 @@ private void transferQueuedEventsIfNecessary(final boolean forced) { private void retryTransfer(final TabletInsertionEvent tabletInsertionEvent) { if (isTabletBatchModeEnabled) { + // A failed detachment may leave this event in a batch while it is also queued. In that case, + // relinquish the queue reference and let the existing batch transfer it once. + if (tabletBatchBuilder != null && tabletBatchBuilder.containsEvent(tabletInsertionEvent)) { + if (tabletInsertionEvent instanceof EnrichedEvent) { + ((EnrichedEvent) tabletInsertionEvent) + .decreaseReferenceCount(IoTDBDataRegionAsyncSink.class.getName(), false); + clearRetryingEventIfNoHandler((EnrichedEvent) tabletInsertionEvent); + } + return; + } try { - tabletBatchBuilder.onEvent(tabletInsertionEvent); - transferBatchedEventsIfNecessary(); + addTabletEventToBatchAndTransferIfNecessary(tabletInsertionEvent); if (tabletInsertionEvent instanceof EnrichedEvent) { ((EnrichedEvent) tabletInsertionEvent) .decreaseReferenceCount(IoTDBDataRegionAsyncSink.class.getName(), false); + clearRetryingEventIfNoHandler((EnrichedEvent) tabletInsertionEvent); } } catch (final Exception e) { addFailureEventToRetryQueue(tabletInsertionEvent, e); @@ -687,10 +803,6 @@ private void retryTransfer(final TabletInsertionEvent tabletInsertionEvent) { addFailureEventToRetryQueue(tabletInsertionEvent, null); } } catch (final Exception e) { - if (tabletInsertionEvent instanceof EnrichedEvent) { - ((EnrichedEvent) tabletInsertionEvent) - .decreaseReferenceCount(IoTDBDataRegionAsyncSink.class.getName(), false); - } addFailureEventToRetryQueue(tabletInsertionEvent, e); } } @@ -708,6 +820,16 @@ private void retryTransfer(final PipeTsFileInsertionEvent tsFileInsertionEvent) } } + private synchronized void clearRetryingEventIfNoHandler(final EnrichedEvent event) { + if (tabletBatchBuilder != null && tabletBatchBuilder.containsEvent(event)) { + return; + } + final Set handlers = retryingEvent2Handlers.get(event); + if (handlers == null || handlers.isEmpty()) { + retryingEvent2Handlers.remove(event); + } + } + /** * Add failure {@link Event} to retry queue. * @@ -715,11 +837,29 @@ private void retryTransfer(final PipeTsFileInsertionEvent tsFileInsertionEvent) */ @SuppressWarnings("java:S899") public void addFailureEventToRetryQueue(final Event event, final Exception e) { - addFailureEventToRetryQueue(event, e, null); + addFailureEventToRetryQueue(event, e, null, null); + } + + /** Reports a failure from a specific handler so stale callbacks can be ignored. */ + public void addFailureEventToRetryQueue( + final Event event, final Exception e, final PipeTransferTrackableHandler sourceHandler) { + addFailureEventToRetryQueue(event, e, null, sourceHandler); } private synchronized void addFailureEventToRetryQueue( final Event event, final Exception e, final Set> failureRecordedPipes) { + addFailureEventToRetryQueue(event, e, failureRecordedPipes, null); + } + + private synchronized void addFailureEventToRetryQueue( + final Event event, + final Exception e, + final Set> failureRecordedPipes, + final PipeTransferTrackableHandler sourceHandler) { + if (event == null) { + return; + } + final PipeResourceFailureType resourceFailureType = PipeStopStrategy.getResourceFailureType(e, null); isConnectionException = @@ -727,15 +867,18 @@ private synchronized void addFailureEventToRetryQueue( if (event instanceof EnrichedEvent) { final EnrichedEvent enrichedEvent = (EnrichedEvent) event; if (enrichedEvent.isReleased()) { + retryingEvent2Handlers.remove(event); return; } if (isDroppedPipe(enrichedEvent)) { + retryingEvent2Handlers.remove(event); enrichedEvent.clearReferenceCount(IoTDBDataRegionAsyncSink.class.getName()); return; } } if (isClosed.get()) { + retryingEvent2Handlers.remove(event); if (event instanceof EnrichedEvent) { ((EnrichedEvent) event).clearReferenceCount(IoTDBDataRegionAsyncSink.class.getName()); } @@ -750,7 +893,24 @@ private synchronized void addFailureEventToRetryQueue( lastRetryFailureMessage = getRetryFailureMessage(e); } - if (resourceFailureType != null && event instanceof EnrichedEvent) { + final Set retryingHandlers = retryingEvent2Handlers.get(event); + if (retryingHandlers != null) { + if (sourceHandler == null) { + retryingEvent2Handlers.remove(event); + } else if (!retryingHandlers.remove(sourceHandler)) { + return; + } else { + retryingEvent2Handlers.remove(event); + } + } + + final boolean alreadyInRetryQueue = retryEvent2ResourceFailureType.containsKey(event); + final PipeResourceFailureType previousResourceFailureType = + retryEvent2ResourceFailureType.get(event); + + if (resourceFailureType != null + && event instanceof EnrichedEvent + && (!alreadyInRetryQueue || previousResourceFailureType != resourceFailureType)) { final EnrichedEvent enrichedEvent = (EnrichedEvent) event; final Pair pipeKey = new Pair<>(enrichedEvent.getPipeName(), enrichedEvent.getCreationTime()); @@ -761,12 +921,17 @@ private synchronized void addFailureEventToRetryQueue( } } - if (resourceFailureType == null) { - retryEvent2ResourceFailureType.remove(event); - } else { - retryEvent2ResourceFailureType.put(event, resourceFailureType); + // A handler and its outer transfer wrapper may report the same failure. Only the first report + // may transfer ownership to the retry queue. + if (alreadyInRetryQueue) { + if (resourceFailureType != null) { + retryEvent2ResourceFailureType.put(event, resourceFailureType); + } + return; } + retryEvent2ResourceFailureType.put(event, resourceFailureType); + if (event instanceof PipeTsFileInsertionEvent) { retryTsFileQueue.offer((PipeTsFileInsertionEvent) event); retryEventQueueEventCounter.increaseEventCount(event); @@ -778,12 +943,6 @@ private synchronized void addFailureEventToRetryQueue( if (LOGGER.isDebugEnabled()) { LOGGER.debug("Added event {} to retry queue.", event); } - - if (isClosed.get()) { - if (event instanceof EnrichedEvent) { - ((EnrichedEvent) event).clearReferenceCount(IoTDBDataRegionAsyncSink.class.getName()); - } - } } /** @@ -797,6 +956,15 @@ public void addFailureEventsToRetryQueue( events.forEach(event -> addFailureEventToRetryQueue(event, e, failureRecordedPipes)); } + public void addFailureEventsToRetryQueue( + final Iterable events, + final Exception e, + final PipeTransferTrackableHandler sourceHandler) { + final Set> failureRecordedPipes = new HashSet<>(); + events.forEach( + event -> addFailureEventToRetryQueue(event, e, failureRecordedPipes, sourceHandler)); + } + static String formatRetryQueueFailureMessage( final int remainingEvents, final int tabletEventCount, @@ -982,6 +1150,7 @@ && isDroppedPipe((EnrichedEvent) event, committerKey)) { ((EnrichedEvent) event).clearReferenceCount(IoTDBDataRegionAsyncSink.class.getName()); retryEventQueueEventCounter.decreaseEventCount(event); retryEvent2ResourceFailureType.remove(event); + retryingEvent2Handlers.remove(event); return true; } return false; @@ -994,20 +1163,34 @@ && isDroppedPipe((EnrichedEvent) event, committerKey)) { ((EnrichedEvent) event).clearReferenceCount(IoTDBDataRegionAsyncSink.class.getName()); retryEventQueueEventCounter.decreaseEventCount(event); retryEvent2ResourceFailureType.remove(event); + retryingEvent2Handlers.remove(event); return true; } return false; }); + retryingEvent2Handlers + .keySet() + .removeIf( + event -> + event instanceof EnrichedEvent + && isDroppedPipe((EnrichedEvent) event, committerKey)); + if (retryEventQueue.isEmpty() && retryTsFileQueue.isEmpty()) { lastRetryFailureMessage = null; } } @Override - // synchronized to avoid close connector when transfer event - public synchronized void close() { - isClosed.set(true); + public void close() { + final Set handlersToClose; + synchronized (this) { + if (!isClosed.compareAndSet(false, true)) { + return; + } + // Handler callbacks can call back into this sink, so close them outside the sink monitor. + handlersToClose = ImmutableSet.copyOf(pendingHandlers.keySet()); + } syncConnector.close(); @@ -1015,15 +1198,11 @@ public synchronized void close() { tabletBatchBuilder.close(); } - // ensure all on-the-fly handlers have been cleared - if (hasPendingHandlers()) { - ImmutableSet.copyOf(pendingHandlers.keySet()) - .forEach( - handler -> { - handler.clearEventsReferenceCount(); - eliminateHandler(handler, true); - }); - } + handlersToClose.forEach( + handler -> { + handler.clearEventsReferenceCount(); + eliminateHandler(handler, true); + }); try { if (clientManager != null) { @@ -1051,12 +1230,14 @@ public synchronized void clearRetryEventsReferenceCount() { retryTsFileQueue.isEmpty() ? retryEventQueue.poll() : retryTsFileQueue.poll(); retryEventQueueEventCounter.decreaseEventCount(event); retryEvent2ResourceFailureType.remove(event); + retryingEvent2Handlers.remove(event); if (event instanceof EnrichedEvent) { ((EnrichedEvent) event).clearReferenceCount(IoTDBDataRegionAsyncSink.class.getName()); } } retryEvent2ResourceFailureType.clear(); lastRetryFailureMessage = null; + retryingEvent2Handlers.clear(); } //////////////////////// APIs provided for metric framework //////////////////////// @@ -1080,7 +1261,40 @@ public boolean isClosed() { } public void trackHandler(final PipeTransferTrackableHandler handler) { - pendingHandlers.put(handler, handler); + boolean closeImmediately = false; + synchronized (this) { + if (isClosed.get()) { + closeImmediately = true; + } else { + pendingHandlers.put(handler, handler); + } + } + + if (closeImmediately) { + handler.clearEventsReferenceCount(); + eliminateHandler(handler, true); + } + } + + /** Registers a newly created handler with any retry attempt currently owning its events. */ + public synchronized void trackRetryHandler( + final PipeTransferTrackableHandler handler, final Iterable events) { + for (final Event event : events) { + final Set handlers = retryingEvent2Handlers.get(event); + if (handlers != null) { + handlers.add(handler); + } + } + } + + private synchronized void untrackRetryHandler(final PipeTransferTrackableHandler handler) { + retryingEvent2Handlers + .entrySet() + .removeIf( + entry -> { + entry.getValue().remove(handler); + return entry.getValue().isEmpty(); + }); } public void eliminateHandler( @@ -1090,6 +1304,7 @@ public void eliminateHandler( } handler.close(); pendingHandlers.remove(handler); + untrackRetryHandler(handler); } public boolean hasPendingHandlers() { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTabletBatchEventHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTabletBatchEventHandler.java index d5ee15c6367b..a535e6e0c1c5 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTabletBatchEventHandler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTabletBatchEventHandler.java @@ -121,7 +121,7 @@ protected void onErrorInternal(final Exception exception) { events.size(), events.stream().map(EnrichedEvent::getPipeName).collect(Collectors.toSet())); } finally { - sink.addFailureEventsToRetryQueue(events, exception); + sink.addFailureEventsToRetryQueue(events, exception, this); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTabletInsertionEventHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTabletInsertionEventHandler.java index a8f1136b897d..e78b19b36bb3 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTabletInsertionEventHandler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTabletInsertionEventHandler.java @@ -108,7 +108,7 @@ protected void onErrorInternal(final Exception exception) { event instanceof EnrichedEvent ? ((EnrichedEvent) event).getCommitterKey() : null, event instanceof EnrichedEvent ? ((EnrichedEvent) event).getCommitIds() : null); } finally { - sink.addFailureEventToRetryQueue(event, exception); + sink.addFailureEventToRetryQueue(event, exception, this); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTrackableHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTrackableHandler.java index b2c914954965..1f3c71ea209c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTrackableHandler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTrackableHandler.java @@ -37,6 +37,7 @@ import org.slf4j.LoggerFactory; import java.util.Objects; +import java.util.concurrent.atomic.AtomicBoolean; public abstract class PipeTransferTrackableHandler implements AsyncMethodCallback, AutoCloseable { @@ -44,24 +45,42 @@ public abstract class PipeTransferTrackableHandler protected final IoTDBDataRegionAsyncSink sink; protected volatile AsyncPipeDataTransferServiceClient client; + private final AtomicBoolean terminal = new AtomicBoolean(false); public PipeTransferTrackableHandler(final IoTDBDataRegionAsyncSink sink) { this.sink = sink; } @Override - public void onComplete(final TPipeTransferResp response) { + public synchronized void onComplete(final TPipeTransferResp response) { + if (terminal.get()) { + return; + } + if (Objects.nonNull(client) && Objects.nonNull(response)) { sink.recordReceiverStatus(client.getEndPoint(), response.getStatus()); } if (sink.isClosed()) { + if (!terminal.compareAndSet(false, true)) { + return; + } clearEventsReferenceCount(); sink.eliminateHandler(this, true); return; } - if (onCompleteInternal(response)) { + final boolean completed; + try { + completed = onCompleteInternal(response); + } catch (final Exception e) { + onError(e); + return; + } + if (completed) { + if (!terminal.compareAndSet(false, true)) { + return; + } // eliminate handler only when all transmissions corresponding to the handler have been // completed // NOTE: We should not clear the reference count of events, as this would cause the @@ -71,10 +90,20 @@ public void onComplete(final TPipeTransferResp response) { } @Override - public void onError(final Exception exception) { + public synchronized void onError(final Exception exception) { + if (!terminal.compareAndSet(false, true)) { + return; + } + if (client != null) { - ThriftClient.resolveException(exception, client); - client.setPrintLogWhenEncounterException(false); + try { + ThriftClient.resolveException(exception, client); + } catch (final Exception resolveException) { + exception.addSuppressed(resolveException); + LOGGER.warn("Failed to resolve transfer exception.", resolveException); + } finally { + client.setPrintLogWhenEncounterException(false); + } } if (sink.isClosed()) { @@ -83,8 +112,11 @@ public void onError(final Exception exception) { return; } - onErrorInternal(exception); - sink.eliminateHandler(this, false); + try { + onErrorInternal(exception); + } finally { + sink.eliminateHandler(this, false); + } } /** @@ -96,15 +128,18 @@ public void onError(final Exception exception) { * is closed or the receiver probe is delayed * @throws TException if an error occurs during the transfer */ - protected boolean tryTransfer( + protected synchronized boolean tryTransfer( final AsyncPipeDataTransferServiceClient client, final TPipeTransferReq req) throws TException { + if (terminal.get()) { + return false; + } if (Objects.isNull(this.client)) { this.client = client; } // track handler before checking if connector is closed sink.trackHandler(this); - if (returnFalseIfSinkIsClosed(client)) { + if (handleSinkClosed(client)) { return false; } try { @@ -114,31 +149,36 @@ protected boolean tryTransfer( onError(e); return false; } - if (returnFalseIfSinkIsClosed(client)) { + if (handleSinkClosed(client)) { return false; } doTransfer(client, req); return true; } - private boolean returnFalseIfSinkIsClosed(final AsyncPipeDataTransferServiceClient client) { + private synchronized boolean handleSinkClosed(final AsyncPipeDataTransferServiceClient client) { if (!sink.isClosed()) { return false; } + if (!terminal.compareAndSet(false, true)) { + return true; + } clearEventsReferenceCount(); sink.eliminateHandler(this, true); - client.setShouldReturnSelf(true); - client.returnSelf( - (e) -> { - if (e instanceof IllegalStateException) { - PipeLogger.log( - LOGGER::info, - "Illegal state when return the client to object pool, maybe the pool is already cleared. Will ignore."); - return true; - } - return false; - }); + if (client != null) { + client.setShouldReturnSelf(true); + client.returnSelf( + (e) -> { + if (e instanceof IllegalStateException) { + PipeLogger.log( + LOGGER::info, + "Illegal state when return the client to object pool, maybe the pool is already cleared. Will ignore."); + return true; + } + return false; + }); + } this.client = null; return true; } @@ -294,7 +334,7 @@ private void fallbackToWholeRequest( try { client.setShouldReturnSelf(shouldReturnSelf); sink.waitIfReceiverRetryIsBackedOff(client.getEndPoint()); - if (returnFalseIfSinkIsClosed(client)) { + if (handleSinkClosed(client)) { return; } client.pipeTransfer(originalReq, this); @@ -324,6 +364,6 @@ public void closeClient() { @Override public void close() { - // Do nothing + terminal.set(true); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandler.java index 8e13e3171c30..b8b757250f8c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandler.java @@ -308,13 +308,12 @@ protected void mayLimitRateAndRecordIO(final long requiredBytes) { } @Override - public void onComplete(final TPipeTransferResp response) { + public synchronized void onComplete(final TPipeTransferResp response) { try { super.onComplete(response); } finally { if (sink.isClosed()) { - releaseReadBufferMemoryBlock(); - returnClientIfNecessary(); + releaseReadBufferAndReturnClient(); } } } @@ -340,21 +339,7 @@ protected boolean onCompleteInternal(final TPipeTransferResp response) { } try { - if (reader != null) { - reader.close(); - } - - // Delete current file when using tsFile as batch - if (events.stream().anyMatch(event -> !(event instanceof PipeTsFileInsertionEvent))) { - RetryUtils.retryOnException( - () -> { - FileUtils.delete(currentFile); - return null; - }); - } - } catch (final IOException e) { - LOGGER.warn( - "Failed to close file reader or delete tsFile when successfully transferred file.", e); + closeReaderAndDeleteBatchFile(true); } finally { final int referenceCount = eventsReferenceCount.decrementAndGet(); if (referenceCount <= 0) { @@ -377,8 +362,7 @@ protected boolean onCompleteInternal(final TPipeTransferResp response) { referenceCount); } - releaseReadBufferMemoryBlock(); - returnClientIfNecessary(); + releaseReadBufferAndReturnClient(); } return true; @@ -416,12 +400,11 @@ protected boolean onCompleteInternal(final TPipeTransferResp response) { } @Override - public void onError(final Exception exception) { + public synchronized void onError(final Exception exception) { try { super.onError(exception); } finally { - releaseReadBufferMemoryBlock(); - returnClientIfNecessary(); + releaseReadBufferAndReturnClient(); } } @@ -456,27 +439,13 @@ protected void onErrorInternal(final Exception exception) { } try { - if (reader != null) { - reader.close(); - } - - // Delete current file when using tsFile as batch - if (events.stream().anyMatch(event -> !(event instanceof PipeTsFileInsertionEvent))) { - RetryUtils.retryOnException( - () -> { - FileUtils.delete(currentFile); - return null; - }); - } - } catch (final IOException e) { - LOGGER.warn("Failed to close file reader or delete tsFile when failed to transfer file.", e); + closeReaderAndDeleteBatchFile(false); } finally { try { - releaseReadBufferMemoryBlock(); - returnClientIfNecessary(); + releaseReadBufferAndReturnClient(); } finally { if (eventsHadBeenAddedToRetryQueue.compareAndSet(false, true)) { - sink.addFailureEventsToRetryQueue(events, exception); + sink.addFailureEventsToRetryQueue(events, exception, this); } } } @@ -527,34 +496,66 @@ public void clearEventsReferenceCount() { } @Override - public void close() { + public synchronized void close() { try { - if (reader != null) { - reader.close(); - reader = null; + closeReaderAndDeleteBatchFile(false); + } finally { + try { + super.close(); + } finally { + releaseReadBufferMemoryBlock(); } + } + } - if (currentFile.exists() - && events.stream().anyMatch(event -> !(event instanceof PipeTsFileInsertionEvent))) { + private void closeReaderAndDeleteBatchFile(final boolean transferSucceeded) { + final String errorMessage = + transferSucceeded + ? "Failed to close file reader or delete tsFile after successful transfer." + : "Failed to close file reader or delete tsFile."; + final RandomAccessFile readerToClose = reader; + if (readerToClose != null) { + try { + RetryUtils.retryOnException( + () -> { + readerToClose.close(); + return null; + }); + if (reader == readerToClose) { + reader = null; + } + } catch (final IOException e) { + LOGGER.warn(errorMessage, e); + } + } + if (currentFile.exists() + && events.stream().anyMatch(event -> !(event instanceof PipeTsFileInsertionEvent))) { + try { RetryUtils.retryOnException( () -> { FileUtils.delete(currentFile); return null; }); + } catch (final IOException e) { + LOGGER.warn(errorMessage, e); } - } catch (final IOException e) { - LOGGER.warn("Failed to close file reader or delete generated batch file.", e); - } finally { - super.close(); + } + } + + private void releaseReadBufferAndReturnClient() { + try { releaseReadBufferMemoryBlock(); + } finally { + returnClientIfNecessary(); } } private void releaseReadBufferMemoryBlock() { - if (memoryBlock != null) { - memoryBlock.close(); - memoryBlock = null; - readBuffer = null; + final PipeTsFileMemoryBlock block = memoryBlock; + memoryBlock = null; + readBuffer = null; + if (block != null) { + block.close(); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java index 728824b4cd90..f9e7d8114bb7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java @@ -21,6 +21,7 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; import org.apache.iotdb.commons.pipe.agent.task.progress.CommitterKey; import org.apache.iotdb.commons.pipe.config.PipeConfig; import org.apache.iotdb.commons.pipe.event.EnrichedEvent; @@ -143,7 +144,17 @@ public void transfer(final TabletInsertionEvent tabletInsertionEvent) throws Exc try { if (isTabletBatchModeEnabled) { - tabletBatchBuilder.onEvent(tabletInsertionEvent); + try { + tabletBatchBuilder.onEvent(tabletInsertionEvent); + } catch (final PipeRuntimeOutOfMemoryCriticalException memoryException) { + try { + doTransferWrapper(); + } catch (final Exception transferException) { + transferException.addSuppressed(memoryException); + throw transferException; + } + tabletBatchBuilder.onEvent(tabletInsertionEvent); + } doTransferWrapper(); } else { if (tabletInsertionEvent instanceof PipeInsertNodeTabletInsertionEvent) { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionDataContainerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionDataContainerTest.java index 1d67bf68732d..8595e3124573 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionDataContainerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionDataContainerTest.java @@ -88,6 +88,8 @@ import java.util.List; import java.util.Objects; import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; @@ -386,6 +388,66 @@ public void testConsumeTabletInsertionEventsWithRetryKeepsParserForTransientOutO event.close(); } + @Test(timeout = 60000) + public void testCloseDefersPendingTabletReleaseUntilConsumerReturns() throws Exception { + final PipeTsFileInsertionEvent event = + createPipeTsFileInsertionEventForRetryTest("nonaligned-consume-close-race.tsfile"); + final CountDownLatch consumerEntered = new CountDownLatch(1); + final CountDownLatch allowConsumerReturn = new CountDownLatch(1); + final AtomicReference parsedEventReference = + new AtomicReference<>(); + final AtomicReference consumerFailure = new AtomicReference<>(); + + final Thread consumerThread = + new Thread( + () -> { + try { + event.consumeTabletInsertionEventsWithRetry( + parsedEvent -> { + parsedEventReference.set(parsedEvent); + consumerEntered.countDown(); + try { + allowConsumerReturn.await(); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + throw new RuntimeException("stop after close race"); + }, + "test"); + } catch (final Throwable t) { + consumerFailure.set(t); + } + }, + "pipe-tsfile-close-race-test"); + consumerThread.start(); + + try { + Assert.assertTrue(consumerEntered.await(10, TimeUnit.SECONDS)); + final PipeRawTabletInsertionEvent parsedEvent = parsedEventReference.get(); + Assert.assertNotNull(parsedEvent); + Assert.assertFalse(parsedEvent.isReleased()); + Assert.assertNotNull(getDataContainer(event).get()); + + event.close(); + + // close() detaches the parser immediately, but the tablet remains valid for the callback + // that was already handed it. + Assert.assertFalse(parsedEvent.isReleased()); + Assert.assertNull(getDataContainer(event).get()); + + allowConsumerReturn.countDown(); + consumerThread.join(10_000); + Assert.assertFalse(consumerThread.isAlive()); + Assert.assertTrue(parsedEvent.isReleased()); + Assert.assertNotNull(consumerFailure.get()); + } finally { + allowConsumerReturn.countDown(); + consumerThread.join(10_000); + event.close(); + } + } + private PipeTsFileInsertionEvent createPipeTsFileInsertionEventForRetryTest(final String fileName) throws Exception { nonalignedTsFile = diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManagerResizeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManagerResizeTest.java index 08b1d8960042..08effc87e86f 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManagerResizeTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManagerResizeTest.java @@ -33,6 +33,7 @@ public class PipeMemoryManagerResizeTest { private static final long TOTAL_MEMORY_SIZE_IN_BYTES = 2000; + private static final long TABLET_MEMORY_SIZE_IN_BYTES = 901; private final CommonConfig config = CommonDescriptor.getInstance().getConfig(); private boolean originalMemoryManagementEnabled; @@ -137,6 +138,48 @@ public void testTabletResizeLeavesMemoryForSinkForwardProgress() { Assert.assertEquals(0, manager.getUsedMemorySizeInBytes()); } + @Test + public void testTryResizeRejectsImmediatelyWithoutChangingAccounting() { + final PipeMemoryManager manager = new PipeMemoryManager(TOTAL_MEMORY_SIZE_IN_BYTES, () -> 0); + final PipeTabletMemoryBlock retainedTablet = + manager.forceAllocateForTabletWithRetry(TABLET_MEMORY_SIZE_IN_BYTES); + final PipeTabletMemoryBlock pendingTablet = manager.forceAllocateForTabletWithRetry(0); + + try { + Assert.assertFalse(manager.tryResize(pendingTablet, 1)); + Assert.assertEquals(0, pendingTablet.getMemoryUsageInBytes()); + Assert.assertEquals(TABLET_MEMORY_SIZE_IN_BYTES, manager.getUsedMemorySizeInBytes()); + Assert.assertEquals(TABLET_MEMORY_SIZE_IN_BYTES, manager.getUsedMemorySizeInBytesOfTablets()); + + manager.release(retainedTablet); + Assert.assertTrue(manager.tryResize(pendingTablet, 1)); + Assert.assertEquals(1, pendingTablet.getMemoryUsageInBytes()); + Assert.assertEquals(1, manager.getUsedMemorySizeInBytesOfTablets()); + } finally { + manager.release(retainedTablet); + manager.release(pendingTablet); + } + + Assert.assertEquals(0, manager.getUsedMemorySizeInBytes()); + } + + @Test + public void testTryResizeRejectsNegativeTargetWithoutChangingAccounting() { + final PipeMemoryManager manager = new PipeMemoryManager(TOTAL_MEMORY_SIZE_IN_BYTES, () -> 0); + final PipeTabletMemoryBlock tablet = manager.forceAllocateForTabletWithRetry(10); + + try { + Assert.assertFalse(manager.tryResize(tablet, -1)); + Assert.assertEquals(10, tablet.getMemoryUsageInBytes()); + Assert.assertEquals(10, manager.getUsedMemorySizeInBytes()); + Assert.assertEquals(10, manager.getUsedMemorySizeInBytesOfTablets()); + } finally { + manager.release(tablet); + } + + Assert.assertEquals(0, manager.getUsedMemorySizeInBytes()); + } + @Test public void testFloatingAndNonFloatingMemoryShareTheSamePool() { final AtomicLong floatingMemoryUsageInBytes = new AtomicLong(0); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeSinkTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeSinkTest.java index 537ae6648b2b..62dca5a762eb 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeSinkTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeSinkTest.java @@ -143,8 +143,11 @@ public void testAsyncSinkDropDoesNotRequeueDroppedPipeEvents() throws Exception recreatedPipeEvent.setCommitterKeyAndCommitId(new CommitterKey("pipe", 2L, 1, -1), 1L); connector.addFailureEventToRetryQueue(recreatedPipeEvent, new PipeException("test")); + connector.addFailureEventToRetryQueue(recreatedPipeEvent, new PipeException("test-again")); Assert.assertEquals(1, connector.getRetryEventQueueSize()); + connector.clearRetryEventsReferenceCount(); + Assert.assertTrue(recreatedPipeEvent.isReleased()); } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilderTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilderTest.java new file mode 100644 index 000000000000..80d284cf0e6b --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilderTest.java @@ -0,0 +1,109 @@ +/* + * 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.db.pipe.sink.payload.evolvable.batch; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; +import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent; +import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters; +import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.utils.Pair; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.Assert; +import org.junit.Test; + +import java.util.Collections; +import java.util.List; + +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_IOTDB_BATCH_DELAY_MS_KEY; + +public class PipeTransferBatchReqBuilderTest { + + @Test + public void testDetachedBatchCannotClearSubsequentEvent() throws Exception { + final PipeTransferBatchReqBuilder builder = + new PipeTransferBatchReqBuilder( + new PipeParameters(Collections.singletonMap(CONNECTOR_IOTDB_BATCH_DELAY_MS_KEY, "0"))); + final PipeRawTabletInsertionEvent firstEvent = createEvent(1); + final PipeRawTabletInsertionEvent secondEvent = createEvent(2); + + try { + builder.onEvent(firstEvent); + final List> detachedBatches = + builder.getAllNonEmptyAndShouldEmitBatchesAndDetach(); + Assert.assertEquals(1, detachedBatches.size()); + + builder.onEvent(secondEvent); + detachedBatches.get(0).getRight().closeAfterEventTransfer(); + + Assert.assertEquals(1, builder.size()); + Assert.assertEquals(1, secondEvent.getReferenceCount()); + Assert.assertFalse(secondEvent.isReleased()); + + // Simulate completion by the handler that owns the detached event reference. + firstEvent.decreaseReferenceCount(getClass().getName(), false); + Assert.assertTrue(firstEvent.isReleased()); + } finally { + builder.close(); + } + + Assert.assertTrue(secondEvent.isReleased()); + } + + @Test + public void testMemoryPressureKeepsExistingBatchEmittable() throws Exception { + final PipeTabletEventBatch batch = + new PipeTabletEventBatch(Integer.MAX_VALUE, Long.MAX_VALUE, null) { + private int constructCount; + + @Override + protected boolean constructBatch(final TabletInsertionEvent event) { + increaseTotalBufferSizeAndUpdateMemoryBlock( + constructCount++ == 0 ? 1 : Long.MAX_VALUE / 2); + return true; + } + }; + + try { + Assert.assertFalse(batch.onEvent(createEvent(1))); + Assert.assertThrows( + PipeRuntimeOutOfMemoryCriticalException.class, () -> batch.onEvent(createEvent(2))); + Assert.assertTrue(batch.shouldEmit()); + } finally { + batch.close(); + } + } + + private static PipeRawTabletInsertionEvent createEvent(final int value) { + final Tablet tablet = + new Tablet( + "root.test.device", + Collections.singletonList(new MeasurementSchema("s1", TSDataType.INT32)), + 1); + tablet.addTimestamp(0, value); + tablet.addValue("s1", 0, value); + tablet.rowSize = 1; + return new PipeRawTabletInsertionEvent( + false, "root.test", null, "root.test", tablet, false, null, 0, null, null, false); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java index fbb985a40b6f..fbf83539217f 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java @@ -133,6 +133,7 @@ private PipeRawTabletInsertionEvent createPipeRawTabletInsertionEvent( final Tablet tablet = new Tablet("root.db.d1", schemaList, 1); tablet.addTimestamp(0, value); tablet.addValue("s1", 0, value); + tablet.rowSize = 1; return new PipeRawTabletInsertionEvent( tablet, false, pipeName, creationTime, null, null, false); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTrackableHandlerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTrackableHandlerTest.java index a0f71536ffff..522a1d06fe5d 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTrackableHandlerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTrackableHandlerTest.java @@ -29,6 +29,7 @@ import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeRequestType; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferSliceReq; import org.apache.iotdb.db.pipe.sink.protocol.thrift.async.IoTDBDataRegionAsyncSink; +import org.apache.iotdb.pipe.api.exception.PipeException; import org.apache.iotdb.rpc.TSStatusCode; import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; import org.apache.iotdb.service.rpc.thrift.TPipeTransferResp; @@ -47,6 +48,8 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; public class PipeTransferTrackableHandlerTest { @@ -220,6 +223,123 @@ public void testClientIsReturnedWhenReceiverProbeIsDelayed() throws Exception { .pipeTransfer(Mockito.any(TPipeTransferReq.class), Mockito.any()); } + @Test + public void testTerminalCallbacksAreIdempotent() { + final IoTDBDataRegionAsyncSink sink = Mockito.mock(IoTDBDataRegionAsyncSink.class); + final TestPipeTransferTrackableHandler completeHandler = + new TestPipeTransferTrackableHandler(sink); + + completeHandler.onComplete(successResp()); + completeHandler.onComplete(successResp()); + + Assert.assertEquals(1, completeHandler.completeCount); + Mockito.verify(sink, Mockito.times(1)).eliminateHandler(completeHandler, false); + + final TestPipeTransferTrackableHandler errorHandler = + new TestPipeTransferTrackableHandler(sink); + errorHandler.onError(new PipeException("first")); + errorHandler.onError(new PipeException("second")); + + Assert.assertEquals(1, errorHandler.errorCount); + Mockito.verify(sink, Mockito.times(1)).eliminateHandler(errorHandler, false); + } + + @Test + public void testCompletionFailureReachesErrorCallback() throws Exception { + commonConfig.setPipeSinkRequestSliceThresholdBytes(1024); + final IoTDBDataRegionAsyncSink sink = Mockito.mock(IoTDBDataRegionAsyncSink.class); + final AsyncPipeDataTransferServiceClient client = + Mockito.mock(AsyncPipeDataTransferServiceClient.class); + final TEndPoint endPoint = new TEndPoint("127.0.0.1", 6667); + final PipeException exception = new PipeException("recording receiver status failed"); + Mockito.when(client.getEndPoint()).thenReturn(endPoint); + Mockito.doThrow(exception) + .when(sink) + .recordReceiverStatus(Mockito.eq(endPoint), Mockito.any(TSStatus.class)); + Mockito.doAnswer( + invocation -> { + final AsyncMethodCallback callback = invocation.getArgument(1); + // TAsyncMethodCall reports exceptions from onComplete through onError. + try { + callback.onComplete(successResp()); + } catch (final Exception e) { + callback.onError(e); + } + return null; + }) + .when(client) + .pipeTransfer(Mockito.any(TPipeTransferReq.class), Mockito.any()); + + final TestPipeTransferTrackableHandler handler = new TestPipeTransferTrackableHandler(sink); + handler.transfer(client, createReq(1)); + + Assert.assertEquals(0, handler.completeCount); + Assert.assertEquals(1, handler.errorCount); + Mockito.verify(sink).eliminateHandler(handler, false); + } + + @Test + public void testSinkCloseDoesNotDeadlockWithHandlerCallback() throws Exception { + final CloseAwareAsyncSink sink = new CloseAwareAsyncSink(); + final CountDownLatch callbackEntered = new CountDownLatch(1); + final CountDownLatch allowCallbackToFinish = new CountDownLatch(1); + final PipeTransferTrackableHandler handler = + new PipeTransferTrackableHandler(sink) { + @Override + protected boolean onCompleteInternal(final TPipeTransferResp response) { + callbackEntered.countDown(); + try { + allowCallbackToFinish.await(); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + return false; + } + return true; + } + + @Override + protected void onErrorInternal(final Exception exception) { + // No-op. + } + + @Override + protected void doTransfer( + final AsyncPipeDataTransferServiceClient client, final TPipeTransferReq req) { + // No-op. + } + + @Override + public void clearEventsReferenceCount() { + // No-op. + } + }; + sink.trackHandler(handler); + + final Thread callbackThread = + new Thread(() -> handler.onComplete(successResp()), "pipe-handler-callback"); + callbackThread.setDaemon(true); + final Thread closeThread = new Thread(sink::close, "pipe-sink-close"); + closeThread.setDaemon(true); + + callbackThread.start(); + Assert.assertTrue(callbackEntered.await(5, TimeUnit.SECONDS)); + closeThread.start(); + Assert.assertTrue(sink.closeEliminationStarted.await(5, TimeUnit.SECONDS)); + + try { + allowCallbackToFinish.countDown(); + callbackThread.join(TimeUnit.SECONDS.toMillis(5)); + closeThread.join(TimeUnit.SECONDS.toMillis(5)); + Assert.assertFalse(callbackThread.isAlive()); + Assert.assertFalse(closeThread.isAlive()); + Assert.assertTrue(sink.isClosed()); + } finally { + allowCallbackToFinish.countDown(); + callbackThread.interrupt(); + closeThread.interrupt(); + } + } + @Test public void testReceiverRetriesAreSerializedForAnyFailureStatus() { commonConfig.setPipeSinkSubtaskSleepIntervalInitMs(40); @@ -330,4 +450,24 @@ public void clearEventsReferenceCount() { // Do nothing } } + + private static class CloseAwareAsyncSink extends IoTDBDataRegionAsyncSink { + private final CountDownLatch closeEliminationStarted = new CountDownLatch(1); + private volatile Thread closeThread; + + @Override + public void close() { + closeThread = Thread.currentThread(); + super.close(); + } + + @Override + public void eliminateHandler( + final PipeTransferTrackableHandler handler, final boolean closeClient) { + if (Thread.currentThread() == closeThread) { + closeEliminationStarted.countDown(); + } + super.eliminateHandler(handler, closeClient); + } + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandlerCleanupTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandlerCleanupTest.java index 892295b34b34..3cca3285cca1 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandlerCleanupTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandlerCleanupTest.java @@ -33,6 +33,8 @@ import org.mockito.Mockito; import java.io.File; +import java.io.IOException; +import java.io.RandomAccessFile; import java.lang.reflect.Field; import java.nio.file.Files; import java.util.Collections; @@ -51,6 +53,23 @@ public void testCloseDeletesBatchFile() throws Exception { Assert.assertFalse(file.exists()); } + @Test + public void testCloseDeletesBatchFileWhenReaderCloseFails() throws Exception { + final File file = Files.createTempFile("pipe-transfer-close-failure", ".tsfile").toFile(); + final EnrichedEvent event = Mockito.mock(EnrichedEvent.class); + final PipeTransferTsFileHandler handler = createHandler(file, event); + final RandomAccessFile reader = Mockito.mock(RandomAccessFile.class); + Mockito.doThrow(new IOException("close failed")).when(reader).close(); + final Field readerField = PipeTransferTsFileHandler.class.getDeclaredField("reader"); + readerField.setAccessible(true); + readerField.set(handler, reader); + + handler.close(); + + Mockito.verify(reader, Mockito.atLeastOnce()).close(); + Assert.assertFalse(file.exists()); + } + @Test public void testNullClientDeletesBatchFile() throws Exception { final File file = Files.createTempFile("pipe-transfer-null-client", ".tsfile").toFile(); @@ -111,7 +130,9 @@ public void testSealFailurePassesNestedReceiverMessageToRetryQueue() throws Exce final ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); Mockito.verify(sink) .addFailureEventsToRetryQueue( - Mockito.eq(Collections.singletonList(event)), exceptionCaptor.capture()); + Mockito.eq(Collections.singletonList(event)), + exceptionCaptor.capture(), + Mockito.eq(handler)); Assert.assertEquals("receiver disk is full", exceptionCaptor.getValue().getMessage()); } finally { if (file.exists()) {