diff --git a/nifi-docs/src/main/asciidoc/user-guide.adoc b/nifi-docs/src/main/asciidoc/user-guide.adoc index b5c0eee2c966..3790d4043d17 100644 --- a/nifi-docs/src/main/asciidoc/user-guide.adoc +++ b/nifi-docs/src/main/asciidoc/user-guide.adoc @@ -1781,6 +1781,14 @@ Stateless Engine is not a safe choice if data must be persisted by NiFi. However it can be a great choice. Additionally, for protocols such as HTTP, NiFi offers processors that are capable of receiving data, performing some processing, and then sending an acknowledgment. As such, it is safe even when NiFi is responsible for accepting incoming connections. The key here is the application-level acknowledgment message that is sent from NiFi. +A Process Group using the Stateless Engine can also be configured with a maximum amount of FlowFile content to keep in memory. This is controlled by two settings: a data size such as `4 GB`, +and a heap percentage greater than `0` and no more than `90`. A value of `0 B` means zero bytes -- it does not mean unlimited. If only one setting is configured, that setting is used. If both are +configured, the smaller of the two limits is used, so `80` percent with a `4 GB` maximum means 80% of the heap, up to 4 GB. If neither is configured, or if the effective limit is zero, +all content is written to the Content Repository. By default, neither setting is configured, which uses no in-memory content. The configured maximum applies to the Stateless Process Group +as a whole, including all Concurrent Tasks. The limit counts allocated content buffer capacity, including unused space in those buffers. Object metadata and buffers allocated by processors +are additional heap usage, so leave sufficient heap available for them. In-memory content is dropped on restart, just like other data inside a Stateless group. Content that has already spilled to the Content Repository, +or that has left the group through an Output Port, remains in the Content Repository. + ==== Data Ordering @@ -1802,7 +1810,8 @@ the data is would still be queued up outside of the Stateless Engine, and this w As such, the Provenance Events are not stored into the Provenance Repository until the transaction completes for a Stateless flow. If a FlowFile is routed to a Failure Port, or if the invocation times out, the Provenance Events are discard. There are, however, two exceptions to this rule: `SEND` and `REMOTE_INVOCATION` events. Even if the transaction is rolled back, the fact that data was sent, or that some remote invocation occurred cannot be rolled back. Therefore, the Provenance Repository is still updated to note the fact that -these events occurred. +these events occurred. When FlowFile content is buffered in memory, those Provenance events still refer to that in-memory content, so the content cannot be downloaded or replayed from +the events after the transaction completes. ==== Site-to-Site diff --git a/nifi-framework-api/src/main/java/org/apache/nifi/controller/repository/TransferableContentClaimRepository.java b/nifi-framework-api/src/main/java/org/apache/nifi/controller/repository/TransferableContentClaimRepository.java new file mode 100644 index 000000000000..96a9bec7c8d2 --- /dev/null +++ b/nifi-framework-api/src/main/java/org/apache/nifi/controller/repository/TransferableContentClaimRepository.java @@ -0,0 +1,43 @@ +/* + * 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.nifi.controller.repository; + +import org.apache.nifi.controller.repository.claim.ContentClaim; + +import java.io.IOException; + +/** + * Content Repository that supports transferring Content Claims to a backing Content Repository. + */ +public interface TransferableContentClaimRepository extends ContentRepository { + + /** + * Prepares a Content Claim from the backing Content Repository. The supplied claim is returned when it already references the backing Content Repository. + * + * @param claim the claim to prepare + * @return a Content Claim whose content is stored in the backing Content Repository + * @throws IOException if the content cannot be written to the backing Content Repository + */ + ContentClaim prepareBackingClaim(ContentClaim claim) throws IOException; + + /** + * Completes the transfer of a prepared Content Claim. + * + * @param claim the original claim passed to {@link #prepareBackingClaim(ContentClaim)} + */ + void completeBackingClaimTransfer(ContentClaim claim); +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-client-dto/src/main/java/org/apache/nifi/web/api/dto/ProcessGroupDTO.java b/nifi-framework-bundle/nifi-framework/nifi-client-dto/src/main/java/org/apache/nifi/web/api/dto/ProcessGroupDTO.java index 17cfa451c51b..741535b32dff 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-client-dto/src/main/java/org/apache/nifi/web/api/dto/ProcessGroupDTO.java +++ b/nifi-framework-bundle/nifi-framework/nifi-client-dto/src/main/java/org/apache/nifi/web/api/dto/ProcessGroupDTO.java @@ -41,6 +41,8 @@ public class ProcessGroupDTO extends ComponentDTO { private String resolvedExecutionEngine; private Integer maxConcurrentTasks; private String statelessFlowTimeout; + private String statelessFlowFileContentInMemoryMax; + private String statelessFlowFileContentInMemoryHeapPercentage; private Integer runningCount; private Integer stoppedCount; @@ -434,4 +436,41 @@ public String getStatelessFlowTimeout() { public void setStatelessFlowTimeout(final String timeout) { this.statelessFlowTimeout = timeout; } + + @Schema(description = "The maximum amount of FlowFile content to buffer in memory when the flow is run using the Stateless Engine, as a data size such as " + + "\"4 GB\". A value of \"0 B\" means zero bytes. When this value is not set, only the heap percentage limit is used. When both this value and the heap percentage are set, " + + "the smaller of the two limits is used.") + public String getStatelessFlowFileContentInMemoryMax() { + return statelessFlowFileContentInMemoryMax; + } + + public void setStatelessFlowFileContentInMemoryMax(final String statelessFlowFileContentInMemoryMax) { + this.statelessFlowFileContentInMemoryMax = statelessFlowFileContentInMemoryMax; + } + + @Schema(description = "The maximum percentage of the Java heap to use for buffering FlowFile content when the flow is run using the Stateless Engine. " + + "The value must be greater than 0 and no more than 90. An empty value means this limit is not configured. When both this value and the in-memory content maximum data size " + + "are set, the smaller of the two limits is used.") + public String getStatelessFlowFileContentInMemoryHeapPercentage() { + return statelessFlowFileContentInMemoryHeapPercentage; + } + + public void setStatelessFlowFileContentInMemoryHeapPercentage(final String statelessFlowFileContentInMemoryHeapPercentage) { + this.statelessFlowFileContentInMemoryHeapPercentage = statelessFlowFileContentInMemoryHeapPercentage; + } + + /** + * @return the heap percentage as an Integer, or null when the value is blank (not configured) + */ + public Integer toStatelessFlowFileContentInMemoryHeapPercentage() { + if (statelessFlowFileContentInMemoryHeapPercentage == null || statelessFlowFileContentInMemoryHeapPercentage.isBlank()) { + return null; + } + + try { + return Integer.valueOf(statelessFlowFileContentInMemoryHeapPercentage.trim()); + } catch (final NumberFormatException e) { + throw new IllegalArgumentException("Max In-Memory Heap Percentage must be a whole number: " + statelessFlowFileContentInMemoryHeapPercentage, e); + } + } } diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/flow/synchronization/StandardVersionedComponentSynchronizer.java b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/flow/synchronization/StandardVersionedComponentSynchronizer.java index e7930d84c329..910a07be88f0 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/flow/synchronization/StandardVersionedComponentSynchronizer.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/flow/synchronization/StandardVersionedComponentSynchronizer.java @@ -501,6 +501,9 @@ private void synchronize(final ProcessGroup group, final VersionedProcessGroup p if (statelessTimeout != null) { group.setStatelessFlowTimeout(statelessTimeout); } + final String statelessFlowFileContentInMemoryMax = proposed.getStatelessFlowFileContentInMemoryMax(); + group.setStatelessContentMaxHeap(statelessFlowFileContentInMemoryMax); + group.setStatelessContentMaxHeapPercentage(proposed.getStatelessFlowFileContentInMemoryHeapPercentage()); if (proposed.getScheduledState() != null && ScheduledState.RUNNING.name().equals(proposed.getScheduledState().name())) { context.getComponentScheduler().startStatelessGroup(group); } diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/groups/StandardProcessGroup.java b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/groups/StandardProcessGroup.java index eb9ae6ee06fc..90c0af28b4da 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/groups/StandardProcessGroup.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/groups/StandardProcessGroup.java @@ -128,6 +128,8 @@ import org.slf4j.LoggerFactory; import java.io.IOException; +import java.math.BigDecimal; +import java.math.RoundingMode; import java.net.ConnectException; import java.nio.charset.StandardCharsets; import java.security.SecureRandom; @@ -138,6 +140,7 @@ import java.util.HashMap; import java.util.HashSet; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.Objects; import java.util.Optional; @@ -156,6 +159,7 @@ import java.util.function.BiFunction; import java.util.function.Function; import java.util.function.Predicate; +import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Collectors; @@ -213,6 +217,9 @@ public final class StandardProcessGroup implements ProcessGroup { private volatile ExecutionEngine executionEngine = ExecutionEngine.INHERITED; private volatile int maxConcurrentTasks = 1; private volatile String statelessFlowTimeout = "1 min"; + private volatile String statelessFlowFileContentInMemoryMax; + private volatile Integer statelessFlowFileContentInMemoryHeapPercentage; + private volatile long resolvedStatelessContentMaxHeap; private volatile Authorizable explicitParentAuthorizable; private final FlowFileActivity flowFileActivity = new ProcessGroupFlowFileActivity(this); @@ -233,6 +240,7 @@ public final class StandardProcessGroup implements ProcessGroup { private static final String DEFAULT_FLOWFILE_EXPIRATION = "0 sec"; private static final long DEFAULT_BACKPRESSURE_OBJECT = 10_000L; private static final String DEFAULT_BACKPRESSURE_DATA_SIZE = "1 GB"; + private static final int MAX_HEAP_PERCENTAGE = 90; private static final Pattern INVALID_DIRECTORY_NAME_CHARACTERS = Pattern.compile("[\\s<>:'\"/\\\\|?*]"); private static final String PATH_SEPARATOR = "/"; private static final String VERSION_SEPARATOR = ":"; @@ -3761,6 +3769,8 @@ private VersionedProcessGroup stripContentsFromRemoteDescendantGroups(final Vers copy.setExecutionEngine(processGroup.getExecutionEngine()); copy.setMaxConcurrentTasks(processGroup.getMaxConcurrentTasks()); copy.setStatelessFlowTimeout(processGroup.getStatelessFlowTimeout()); + copy.setStatelessFlowFileContentInMemoryMax(processGroup.getStatelessFlowFileContentInMemoryMax()); + copy.setStatelessFlowFileContentInMemoryHeapPercentage(processGroup.getStatelessFlowFileContentInMemoryHeapPercentage()); final Set copyChildren = new HashSet<>(); @@ -3785,6 +3795,8 @@ private VersionedProcessGroup stripContentsFromRemoteDescendantGroups(final Vers childCopy.setExecutionEngine(childGroup.getExecutionEngine()); childCopy.setMaxConcurrentTasks(childGroup.getMaxConcurrentTasks()); childCopy.setStatelessFlowTimeout(childGroup.getStatelessFlowTimeout()); + childCopy.setStatelessFlowFileContentInMemoryMax(childGroup.getStatelessFlowFileContentInMemoryMax()); + childCopy.setStatelessFlowFileContentInMemoryHeapPercentage(childGroup.getStatelessFlowFileContentInMemoryHeapPercentage()); copyChildren.add(childCopy); } @@ -4778,6 +4790,132 @@ public void setStatelessFlowTimeout(final String statelessFlowTimeout) { } } + @Override + public String getStatelessContentMaxHeap() { + return statelessFlowFileContentInMemoryMax; + } + + @Override + public void setStatelessContentMaxHeap(final String maxSize) { + writeLock.lock(); + try { + final long proposedMaxSizeBytes = resolveStatelessContentMaxHeap(maxSize, getStatelessContentMaxHeapPercentage()); + verifyCanSetStatelessContentMaxHeap(proposedMaxSizeBytes); + this.statelessFlowFileContentInMemoryMax = normalizeStatelessContentMaxHeap(maxSize); + this.resolvedStatelessContentMaxHeap = proposedMaxSizeBytes; + } finally { + writeLock.unlock(); + } + } + + @Override + public Integer getStatelessContentMaxHeapPercentage() { + return statelessFlowFileContentInMemoryHeapPercentage; + } + + @Override + public void setStatelessContentMaxHeapPercentage(final Integer heapPercentage) { + writeLock.lock(); + try { + final long proposedMaxSizeBytes = resolveStatelessContentMaxHeap(getStatelessContentMaxHeap(), heapPercentage); + verifyCanSetStatelessContentMaxHeap(proposedMaxSizeBytes); + this.statelessFlowFileContentInMemoryHeapPercentage = heapPercentage; + this.resolvedStatelessContentMaxHeap = proposedMaxSizeBytes; + } finally { + writeLock.unlock(); + } + } + + @Override + public long resolveStatelessContentMaxHeap() { + return resolvedStatelessContentMaxHeap; + } + + @Override + public void verifyCanSetStatelessContentMaxHeap(final String maxSize) { + final long proposedMaxSizeBytes = resolveStatelessContentMaxHeap(maxSize, getStatelessContentMaxHeapPercentage()); + verifyCanSetStatelessContentMaxHeap(proposedMaxSizeBytes); + } + + @Override + public void verifyCanSetStatelessContentMaxHeapPercentage(final Integer heapPercentage) { + final long proposedMaxSizeBytes = resolveStatelessContentMaxHeap(getStatelessContentMaxHeap(), heapPercentage); + verifyCanSetStatelessContentMaxHeap(proposedMaxSizeBytes); + } + + private void verifyCanSetStatelessContentMaxHeap(final long proposedMaxSizeBytes) { + // The Content Repository is selected when the Stateless flow starts, so the setting cannot change while the flow is running. + final ProcessGroup statelessGroup = getStatelessGroup(this); + if (statelessGroup != null && statelessGroup.getStatelessScheduledState() != StatelessGroupScheduledState.STOPPED + && proposedMaxSizeBytes != resolveStatelessContentMaxHeap()) { + throw new IllegalStateException("Cannot change the maximum in-memory FlowFile content for " + this + + " while the Stateless flow is running. Stop the Process Group before changing this setting."); + } + } + + private static String normalizeStatelessContentMaxHeap(final String maxSize) { + if (maxSize == null || maxSize.isBlank()) { + return null; + } + + return maxSize.trim(); + } + + private static long resolveStatelessContentMaxHeap(final String maxSize, final Integer heapPercentage) { + validateHeapPercentage(heapPercentage); + + final boolean sizeConfigured = maxSize != null && !maxSize.isBlank(); + final boolean heapPercentageConfigured = heapPercentage != null; + if (!sizeConfigured && !heapPercentageConfigured) { + return 0L; + } + + if (sizeConfigured && !heapPercentageConfigured) { + return parseStatelessContentMaxHeap(maxSize); + } + if (!sizeConfigured) { + return toHeapPercentageBytes(heapPercentage); + } + + return Math.min(parseStatelessContentMaxHeap(maxSize), toHeapPercentageBytes(heapPercentage)); + } + + private static void validateHeapPercentage(final Integer heapPercentage) { + if (heapPercentage != null && (heapPercentage <= 0 || heapPercentage > MAX_HEAP_PERCENTAGE)) { + throw new IllegalArgumentException("Heap percentage must be greater than 0 and no more than " + MAX_HEAP_PERCENTAGE + ": " + heapPercentage); + } + } + + private static long toHeapPercentageBytes(final Integer heapPercentage) { + return BigDecimal.valueOf(heapPercentage) + .multiply(BigDecimal.valueOf(Runtime.getRuntime().maxMemory())) + .divide(BigDecimal.valueOf(100), 0, RoundingMode.DOWN) + .longValue(); + } + + private static long parseStatelessContentMaxHeap(final String maxSize) { + final String normalizedMaxSize = maxSize.trim().toUpperCase(Locale.ROOT); + final Matcher matcher = DataUnit.DATA_SIZE_PATTERN.matcher(normalizedMaxSize); + if (!matcher.matches()) { + throw new IllegalArgumentException("Invalid data size: " + maxSize); + } + + final long multiplier = switch (matcher.group(2)) { + case "B" -> 1L; + case "KB" -> 1L << 10; + case "MB" -> 1L << 20; + case "GB" -> 1L << 30; + case "TB" -> 1L << 40; + default -> throw new IllegalArgumentException("Invalid data size: " + maxSize); + }; + + try { + return new BigDecimal(matcher.group(1)).multiply(BigDecimal.valueOf(multiplier)).longValueExact(); + } catch (final ArithmeticException e) { + throw new IllegalArgumentException("Data size must represent a whole number of bytes within the supported range: " + maxSize, e); + } + } + private void setLoggingAttributes(final boolean recursive) { final Map attributes = new HashMap<>(); diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/registry/flow/mapping/VersionedComponentFlowMapper.java b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/registry/flow/mapping/VersionedComponentFlowMapper.java index 69ac2ff0af06..361e3eabb1b1 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/registry/flow/mapping/VersionedComponentFlowMapper.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/registry/flow/mapping/VersionedComponentFlowMapper.java @@ -274,6 +274,8 @@ private InstantiatedVersionedProcessGroup mapGroup(final ProcessGroup group, fin versionedGroup.setScheduledState(flowMappingOptions.getStateLookup().getState(group)); versionedGroup.setMaxConcurrentTasks(group.getMaxConcurrentTasks()); versionedGroup.setStatelessFlowTimeout(group.getStatelessFlowTimeout()); + versionedGroup.setStatelessFlowFileContentInMemoryMax(group.getStatelessContentMaxHeap()); + versionedGroup.setStatelessFlowFileContentInMemoryHeapPercentage(group.getStatelessContentMaxHeapPercentage()); final ParameterContext parameterContext = group.getParameterContext(); versionedGroup.setParameterContextName(parameterContext == null ? null : parameterContext.getName()); diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/test/java/org/apache/nifi/groups/StandardProcessGroupTest.java b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/test/java/org/apache/nifi/groups/StandardProcessGroupTest.java index 2e2659009673..35829aba44f1 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/test/java/org/apache/nifi/groups/StandardProcessGroupTest.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/test/java/org/apache/nifi/groups/StandardProcessGroupTest.java @@ -28,6 +28,7 @@ import org.apache.nifi.controller.NodeTypeProvider; import org.apache.nifi.controller.ProcessScheduler; import org.apache.nifi.controller.ReloadComponent; +import org.apache.nifi.controller.ScheduledState; import org.apache.nifi.controller.flow.FlowManager; import org.apache.nifi.controller.queue.DropFlowFileRequest; import org.apache.nifi.controller.queue.DropFlowFileState; @@ -48,6 +49,8 @@ import org.mockito.junit.jupiter.MockitoExtension; import java.io.IOException; +import java.math.BigDecimal; +import java.math.RoundingMode; import java.util.Map; import java.util.Optional; import java.util.concurrent.CompletableFuture; @@ -56,6 +59,7 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; @@ -476,6 +480,120 @@ void testDropAllFlowFilesCompletionWaitsForEveryConnection() { assertTrue(completionFuture.isDone()); } + @Test + void testStatelessContentMaxHeapDefaultsToZeroBytesWithUnsetLimits() { + assertNull(processGroup.getStatelessContentMaxHeap()); + assertNull(processGroup.getStatelessContentMaxHeapPercentage()); + assertEquals(0L, processGroup.resolveStatelessContentMaxHeap()); + } + + @Test + void testSetStatelessContentMaxHeapParsesDataSize() { + processGroup.setStatelessContentMaxHeapPercentage(null); + processGroup.setStatelessContentMaxHeap("100 MB"); + assertEquals("100 MB", processGroup.getStatelessContentMaxHeap()); + assertEquals(100L * 1024 * 1024, processGroup.resolveStatelessContentMaxHeap()); + + processGroup.setStatelessContentMaxHeap("1 KB"); + assertEquals(1024L, processGroup.resolveStatelessContentMaxHeap()); + + processGroup.setStatelessContentMaxHeap("1.5 kb"); + assertEquals(1536L, processGroup.resolveStatelessContentMaxHeap()); + + processGroup.setStatelessContentMaxHeap("0 B"); + assertEquals("0 B", processGroup.getStatelessContentMaxHeap()); + assertEquals(0L, processGroup.resolveStatelessContentMaxHeap()); + + processGroup.setStatelessContentMaxHeap(" "); + assertNull(processGroup.getStatelessContentMaxHeap()); + assertEquals(0L, processGroup.resolveStatelessContentMaxHeap()); + } + + @Test + void testSetStatelessContentMaxHeapPercentage() { + processGroup.setStatelessContentMaxHeapPercentage(50); + assertEquals(50, processGroup.getStatelessContentMaxHeapPercentage()); + assertEquals(Runtime.getRuntime().maxMemory() / 2, processGroup.resolveStatelessContentMaxHeap()); + + processGroup.setStatelessContentMaxHeapPercentage(90); + final long expectedNinetyPercent = new BigDecimal("90") + .multiply(BigDecimal.valueOf(Runtime.getRuntime().maxMemory())) + .divide(BigDecimal.valueOf(100), 0, RoundingMode.DOWN) + .longValue(); + assertEquals(expectedNinetyPercent, processGroup.resolveStatelessContentMaxHeap()); + } + + @Test + void testResolveStatelessContentMaxHeapUsesTheSmallerLimitWhenBothAreSet() { + processGroup.setStatelessContentMaxHeap("1 B"); + processGroup.setStatelessContentMaxHeapPercentage(90); + assertEquals(1L, processGroup.resolveStatelessContentMaxHeap()); + + processGroup.setStatelessContentMaxHeap(Runtime.getRuntime().maxMemory() + " B"); + processGroup.setStatelessContentMaxHeapPercentage(50); + assertEquals(Runtime.getRuntime().maxMemory() / 2, processGroup.resolveStatelessContentMaxHeap()); + } + + @Test + void testSetStatelessContentMaxHeapRejectsInvalidLimits() { + processGroup.setStatelessContentMaxHeapPercentage(null); + processGroup.setStatelessContentMaxHeap("1 MB"); + assertThrows(IllegalArgumentException.class, () -> processGroup.setStatelessContentMaxHeap("not a size")); + assertThrows(IllegalArgumentException.class, () -> processGroup.setStatelessContentMaxHeap("-1 MB")); + assertThrows(IllegalArgumentException.class, () -> processGroup.setStatelessContentMaxHeap("limit 1 MB")); + assertThrows(IllegalArgumentException.class, () -> processGroup.setStatelessContentMaxHeap("0.9 B")); + assertThrows(IllegalArgumentException.class, () -> processGroup.setStatelessContentMaxHeap("999999999999999999999999999999999999999999999 TB")); + assertThrows(IllegalArgumentException.class, () -> processGroup.setStatelessContentMaxHeap("50%")); + assertThrows(IllegalArgumentException.class, () -> processGroup.setStatelessContentMaxHeapPercentage(0)); + assertThrows(IllegalArgumentException.class, () -> processGroup.setStatelessContentMaxHeapPercentage(91)); + assertThrows(IllegalArgumentException.class, () -> processGroup.setStatelessContentMaxHeapPercentage(-1)); + assertThrows(IllegalArgumentException.class, () -> processGroup.verifyCanSetStatelessContentMaxHeap("not a size")); + assertThrows(IllegalArgumentException.class, () -> processGroup.verifyCanSetStatelessContentMaxHeapPercentage(0)); + assertEquals("1 MB", processGroup.getStatelessContentMaxHeap()); + assertNull(processGroup.getStatelessContentMaxHeapPercentage()); + assertEquals(1024L * 1024L, processGroup.resolveStatelessContentMaxHeap()); + } + + @Test + void testSetStatelessContentMaxHeapWhileRunningAllowsSameByteCount() { + final StatelessGroupNode statelessGroupNode = mock(StatelessGroupNode.class); + when(statelessGroupNodeFactory.createStatelessGroupNode(any())).thenReturn(statelessGroupNode); + final StandardProcessGroup runningGroup = createStandardProcessGroup("running"); + runningGroup.setStatelessContentMaxHeapPercentage(null); + runningGroup.setStatelessContentMaxHeap("1 MB"); + runningGroup.setExecutionEngine(ExecutionEngine.STATELESS); + when(statelessGroupNode.getCurrentState()).thenReturn(ScheduledState.RUNNING); + + runningGroup.verifyCanSetStatelessContentMaxHeap("1024 KB"); + runningGroup.setStatelessContentMaxHeap("1024 KB"); + assertEquals("1024 KB", runningGroup.getStatelessContentMaxHeap()); + assertThrows(IllegalStateException.class, () -> runningGroup.verifyCanSetStatelessContentMaxHeap("2 MB")); + assertThrows(IllegalStateException.class, () -> runningGroup.setStatelessContentMaxHeap("2 MB")); + assertEquals(1024L * 1024L, runningGroup.resolveStatelessContentMaxHeap()); + + final StandardProcessGroup childGroup = createStandardProcessGroup("child"); + childGroup.setName("Child"); + runningGroup.addProcessGroup(childGroup); + assertThrows(IllegalStateException.class, () -> childGroup.verifyCanSetStatelessContentMaxHeap("2 MB")); + assertThrows(IllegalStateException.class, () -> childGroup.setStatelessContentMaxHeap("2 MB")); + + when(statelessGroupNode.getCurrentState()).thenReturn(ScheduledState.STOPPED); + runningGroup.setStatelessContentMaxHeap("2 MB"); + assertEquals(2L * 1024L * 1024L, runningGroup.resolveStatelessContentMaxHeap()); + runningGroup.setStatelessContentMaxHeap(null); + runningGroup.setStatelessContentMaxHeapPercentage(90); + + when(statelessGroupNode.getCurrentState()).thenReturn(ScheduledState.RUNNING); + runningGroup.verifyCanSetStatelessContentMaxHeapPercentage(90); + runningGroup.setStatelessContentMaxHeapPercentage(90); + assertThrows(IllegalStateException.class, () -> runningGroup.verifyCanSetStatelessContentMaxHeapPercentage(50)); + assertThrows(IllegalStateException.class, () -> runningGroup.setStatelessContentMaxHeapPercentage(50)); + + when(statelessGroupNode.getCurrentState()).thenReturn(ScheduledState.STOPPED); + runningGroup.setStatelessContentMaxHeapPercentage(50); + assertEquals(50, runningGroup.getStatelessContentMaxHeapPercentage()); + } + private StandardProcessGroup createStandardProcessGroup(final String id) { return createStandardProcessGroup(id, null); } diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core-api/src/main/java/org/apache/nifi/groups/ProcessGroup.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core-api/src/main/java/org/apache/nifi/groups/ProcessGroup.java index c96f40c2253a..145fa71a2b89 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-core-api/src/main/java/org/apache/nifi/groups/ProcessGroup.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core-api/src/main/java/org/apache/nifi/groups/ProcessGroup.java @@ -1326,6 +1326,53 @@ default void setConnectorLoggingAttributes(final Map attributes) */ String getStatelessFlowTimeout(); + /** + * @return the configured maximum amount of FlowFile content to buffer in memory when this Process Group is run using the Stateless Execution Engine, + * specified as a data size such as "4 GB". A value of "0 B" means zero bytes. A null or blank value means this limit is not configured. + */ + String getStatelessContentMaxHeap(); + + /** + * Sets the maximum amount of FlowFile content to buffer in memory when this Process Group is run using the Stateless Execution Engine + * @param maxSize the maximum amount of FlowFile content to buffer in memory, as a data size such as "4 GB". A blank value means this limit is not configured. + */ + void setStatelessContentMaxHeap(String maxSize); + + /** + * @return the configured maximum percentage of the Java heap to use for buffering FlowFile content when this Process Group is run using the Stateless + * Execution Engine. The value is greater than 0 and no more than 90. A null value means this limit is not configured. + */ + Integer getStatelessContentMaxHeapPercentage(); + + /** + * Sets the maximum percentage of the Java heap to use for buffering FlowFile content when this Process Group is run using the Stateless Execution Engine + * @param heapPercentage the maximum heap percentage, greater than 0 and no more than 90. A null value means this limit is not configured. + */ + void setStatelessContentMaxHeapPercentage(Integer heapPercentage); + + /** + * @return the amount of FlowFile content that may be buffered in memory, in bytes, when this Process Group is run using the Stateless Execution Engine. + * If only the data size is configured, that size is used. If only the heap percentage is configured, that percentage of the JVM's maximum heap is used. + * If both are configured, the smaller of the two limits is used. If neither is configured, the result is 0, so all FlowFile content is written to the Content Repository. + */ + long resolveStatelessContentMaxHeap(); + + /** + * Verifies that the maximum in-memory FlowFile content data size can be set to the given value. + * @param maxSize the maximum amount of FlowFile content to buffer in memory, as a data size such as "4 GB". A blank value means this limit is not configured. + * @throws IllegalArgumentException if the value is not a valid data size + * @throws IllegalStateException if the value cannot be set because the Stateless flow is running + */ + void verifyCanSetStatelessContentMaxHeap(String maxSize); + + /** + * Verifies that the maximum in-memory FlowFile content heap percentage can be set to the given value. + * @param heapPercentage the maximum heap percentage, greater than 0 and no more than 90. A null value means this limit is not configured. + * @throws IllegalArgumentException if the value is outside the allowed range + * @throws IllegalStateException if the value cannot be set because the Stateless flow is running + */ + void verifyCanSetStatelessContentMaxHeapPercentage(Integer heapPercentage); + /** * @return the FlowFileActivity for this Process Group */ diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/StandardFlowSnippet.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/StandardFlowSnippet.java index 34f6ad108927..c06e95016eb3 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/StandardFlowSnippet.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/StandardFlowSnippet.java @@ -484,6 +484,10 @@ public void instantiate(final FlowManager flowManager, final FlowController flow childGroup.setName(groupDTO.getName()); childGroup.setExecutionEngine(ExecutionEngine.valueOf(groupDTO.getExecutionEngine())); childGroup.setStatelessFlowTimeout(groupDTO.getStatelessFlowTimeout()); + childGroup.setStatelessContentMaxHeap(groupDTO.getStatelessFlowFileContentInMemoryMax()); + if (groupDTO.getStatelessFlowFileContentInMemoryHeapPercentage() != null) { + childGroup.setStatelessContentMaxHeapPercentage(groupDTO.toStatelessFlowFileContentInMemoryHeapPercentage()); + } childGroup.setMaxConcurrentTasks(groupDTO.getMaxConcurrentTasks()); final String flowfileConcurrentName = groupDTO.getFlowfileConcurrency(); diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/flow/StandardStatelessGroupNodeFactory.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/flow/StandardStatelessGroupNodeFactory.java index 6cba62b68a5c..ba5750337ad9 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/flow/StandardStatelessGroupNodeFactory.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/flow/StandardStatelessGroupNodeFactory.java @@ -23,6 +23,7 @@ import org.apache.nifi.controller.FlowController; import org.apache.nifi.controller.kerberos.KerberosConfig; import org.apache.nifi.controller.repository.ContentRepository; +import org.apache.nifi.controller.repository.DeferredStatelessContentRepository; import org.apache.nifi.controller.repository.FlowFileEventRepository; import org.apache.nifi.controller.repository.FlowFileRepository; import org.apache.nifi.controller.repository.NonPurgeableContentRepository; @@ -33,6 +34,7 @@ import org.apache.nifi.controller.scheduling.StatelessProcessScheduler; import org.apache.nifi.controller.scheduling.StatelessProcessSchedulerInitializationContext; import org.apache.nifi.engine.FlowEngine; +import org.apache.nifi.events.EventReporter; import org.apache.nifi.extensions.BundleAvailability; import org.apache.nifi.extensions.ExtensionRepository; import org.apache.nifi.flow.ExternalControllerServiceReference; @@ -122,7 +124,12 @@ public StatelessGroupNode createStatelessGroupNode(final ProcessGroup group) { flowFileRepository.initialize(resourceClaimManager); - final ContentRepository contentRepository = new NonPurgeableContentRepository(flowController.getRepositoryContextFactory().getContentRepository()); + // Defer the choice of Content Repository until it is first used (i.e., when the group starts), because at construction time the group's maximum + // in-memory FlowFile content size has not yet been configured. When that size is greater than zero, content is buffered in memory and spills to the + // NiFi Content Repository once the size is exceeded; otherwise the NiFi instance's Content Repository is used directly. In either case the NiFi + // Content Repository is wrapped so the Stateless flow does not purge content the framework is responsible for cleaning up. + final ContentRepository frameworkContentRepository = new NonPurgeableContentRepository(flowController.getRepositoryContextFactory().getContentRepository()); + final ContentRepository contentRepository = new DeferredStatelessContentRepository(group, frameworkContentRepository, underlyingFlowFileRepository, resourceClaimManager, EventReporter.NO_OP); final RepositoryContextFactory statelessRepoContextFactory = new StatelessRepositoryContextFactory( contentRepository, flowFileRepository, diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/repository/DeferredStatelessContentRepository.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/repository/DeferredStatelessContentRepository.java new file mode 100644 index 000000000000..564a66de9ba8 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/repository/DeferredStatelessContentRepository.java @@ -0,0 +1,262 @@ +/* + * 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.nifi.controller.repository; + +import org.apache.nifi.controller.repository.claim.ContentClaim; +import org.apache.nifi.controller.repository.claim.ResourceClaim; +import org.apache.nifi.controller.repository.claim.ResourceClaimManager; +import org.apache.nifi.events.EventReporter; +import org.apache.nifi.groups.ProcessGroup; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.file.Path; +import java.util.Set; + +/** + * A {@link ContentRepository} for an embedded Stateless Process Group that defers the choice of backing repository until it is first used. This is necessary + * because the {@link org.apache.nifi.groups.StatelessGroupNode} is created when the Process Group is constructed, which is before the Process Group's maximum + * in-memory FlowFile content size has been configured. Resolving the backing repository lazily ensures the configured value is honored: when the maximum + * in-memory FlowFile content size is greater than zero, a {@link SpillableContentRepository} buffers FlowFile content in memory up to that size and spills to + * the NiFi instance's Content Repository once it is exceeded; when the size is zero, the NiFi instance's Content Repository is used directly. + */ +public class DeferredStatelessContentRepository implements TransferableContentClaimRepository { + private final ProcessGroup processGroup; + private final ContentRepository contentRepositoryDelegate; + private final FlowFileRepository nifiFlowFileRepository; + private final ResourceClaimManager resourceClaimManager; + private final EventReporter eventReporter; + + private volatile ContentRepository delegate; + private volatile long resolvedMemoryThresholdBytes = Long.MIN_VALUE; + + public DeferredStatelessContentRepository(final ProcessGroup processGroup, final ContentRepository contentRepositoryDelegate, final FlowFileRepository nifiFlowFileRepository, + final ResourceClaimManager resourceClaimManager, final EventReporter eventReporter) { + this.processGroup = processGroup; + this.contentRepositoryDelegate = contentRepositoryDelegate; + this.nifiFlowFileRepository = nifiFlowFileRepository; + this.resourceClaimManager = resourceClaimManager; + this.eventReporter = eventReporter; + } + + /** + * Prepares a claim from the backing Content Repository when the configured delegate buffers content in memory. Otherwise, returns the supplied claim unchanged. + * + * @param claim the claim to prepare + * @return a Content Claim whose content is stored in the backing Content Repository + * @throws IOException if the content cannot be written to the backing Content Repository + */ + @Override + public ContentClaim prepareBackingClaim(final ContentClaim claim) throws IOException { + final ContentRepository resolved = getDelegate(); + if (resolved instanceof final SpillableContentRepository spillableContentRepository) { + return spillableContentRepository.prepareBackingClaim(claim); + } + + return claim; + } + + /** + * Completes the handoff of a prepared Content Claim after the NiFi FlowFile Repository has been updated successfully. + * + * @param claim the original claim passed to {@link #prepareBackingClaim(ContentClaim)} + */ + @Override + public void completeBackingClaimTransfer(final ContentClaim claim) { + final ContentRepository resolved = getDelegate(); + if (resolved instanceof final SpillableContentRepository spillableContentRepository) { + spillableContentRepository.completeBackingClaimTransfer(claim); + } + } + + @Override + public void initialize(final ContentRepositoryContext context) { + // The backing repository is initialized when it is resolved; nothing to do here. + } + + @Override + public void shutdown() { + final ContentRepository resolved = delegate; + if (resolved != null) { + resolved.shutdown(); + } + } + + @Override + public void purge() { + final ContentRepository resolved = delegate; + if (resolved != null) { + resolved.purge(); + } + } + + @Override + public void cleanup() { + final ContentRepository resolved = delegate; + if (resolved != null) { + resolved.cleanup(); + } + } + + @Override + public Set getContainerNames() { + return getDelegate().getContainerNames(); + } + + @Override + public long getContainerCapacity(final String containerName) throws IOException { + return getDelegate().getContainerCapacity(containerName); + } + + @Override + public long getContainerUsableSpace(final String containerName) throws IOException { + return getDelegate().getContainerUsableSpace(containerName); + } + + @Override + public String getContainerFileStoreName(final String containerName) { + return getDelegate().getContainerFileStoreName(containerName); + } + + @Override + public ContentClaim create(final boolean lossTolerant) throws IOException { + return getDelegate().create(lossTolerant); + } + + @Override + public int incrementClaimaintCount(final ContentClaim claim) { + return getDelegate().incrementClaimaintCount(claim); + } + + @Override + public int getClaimantCount(final ContentClaim claim) { + return getDelegate().getClaimantCount(claim); + } + + @Override + public int decrementClaimantCount(final ContentClaim claim) { + return getDelegate().decrementClaimantCount(claim); + } + + @Override + public boolean remove(final ContentClaim claim) { + return getDelegate().remove(claim); + } + + @Override + public ContentClaim clone(final ContentClaim original, final boolean lossTolerant) throws IOException { + return getDelegate().clone(original, lossTolerant); + } + + @Override + public long importFrom(final Path content, final ContentClaim claim) throws IOException { + return getDelegate().importFrom(content, claim); + } + + @Override + public long importFrom(final InputStream content, final ContentClaim claim) throws IOException { + return getDelegate().importFrom(content, claim); + } + + @Override + public long exportTo(final ContentClaim claim, final Path destination, final boolean append) throws IOException { + return getDelegate().exportTo(claim, destination, append); + } + + @Override + public long exportTo(final ContentClaim claim, final Path destination, final boolean append, final long offset, final long length) throws IOException { + return getDelegate().exportTo(claim, destination, append, offset, length); + } + + @Override + public long exportTo(final ContentClaim claim, final OutputStream destination) throws IOException { + return getDelegate().exportTo(claim, destination); + } + + @Override + public long exportTo(final ContentClaim claim, final OutputStream destination, final long offset, final long length) throws IOException { + return getDelegate().exportTo(claim, destination, offset, length); + } + + @Override + public long size(final ContentClaim claim) throws IOException { + return getDelegate().size(claim); + } + + @Override + public long size(final ResourceClaim claim) throws IOException { + return getDelegate().size(claim); + } + + @Override + public InputStream read(final ContentClaim claim) throws IOException { + return getDelegate().read(claim); + } + + @Override + public InputStream read(final ResourceClaim claim) throws IOException { + return getDelegate().read(claim); + } + + @Override + public boolean isResourceClaimStreamSupported() { + return getDelegate().isResourceClaimStreamSupported(); + } + + @Override + public OutputStream write(final ContentClaim claim) throws IOException { + return getDelegate().write(claim); + } + + @Override + public boolean isAccessible(final ContentClaim contentClaim) throws IOException { + return getDelegate().isAccessible(contentClaim); + } + + private ContentRepository getDelegate() { + long memoryThresholdBytes = processGroup.resolveStatelessContentMaxHeap(); + ContentRepository resolved = delegate; + if (resolved != null && resolvedMemoryThresholdBytes == memoryThresholdBytes) { + return resolved; + } + + synchronized (this) { + memoryThresholdBytes = processGroup.resolveStatelessContentMaxHeap(); + resolved = delegate; + if (resolved == null || resolvedMemoryThresholdBytes != memoryThresholdBytes) { + if (resolved != null) { + resolved.shutdown(); + } + + if (memoryThresholdBytes > 0) { + final SpillableContentRepository spillableContentRepository = new SpillableContentRepository(contentRepositoryDelegate, nifiFlowFileRepository, memoryThresholdBytes); + spillableContentRepository.initialize(new StandardContentRepositoryContext(resourceClaimManager, eventReporter)); + resolved = spillableContentRepository; + } else { + resolved = contentRepositoryDelegate; + } + + resolvedMemoryThresholdBytes = memoryThresholdBytes; + delegate = resolved; + } + + return resolved; + } + } +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/repository/SpillableContentRepository.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/repository/SpillableContentRepository.java new file mode 100644 index 000000000000..bd43eeffe5a1 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/repository/SpillableContentRepository.java @@ -0,0 +1,931 @@ +/* + * 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.nifi.controller.repository; + +import org.apache.nifi.controller.repository.claim.ContentClaim; +import org.apache.nifi.controller.repository.claim.ResourceClaim; +import org.apache.nifi.controller.repository.claim.ResourceClaimManager; +import org.apache.nifi.controller.repository.claim.StandardResourceClaim; +import org.apache.nifi.stream.io.StreamUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.io.SequenceInputStream; +import java.nio.file.Files; +import java.nio.file.OpenOption; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.Enumeration; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicLong; + +/** + * A {@link ContentRepository} used by an embedded Stateless Process Group that buffers FlowFile content in memory up to a configured total size and spills to a + * backing Content Repository once that size is exceeded. Content for a single {@link ContentClaim} is stored either entirely in memory or entirely in + * the backing repository: while a claim is being written, each allocation reserves its buffer capacity against the configured size. If capacity is unavailable, + * the bytes buffered so far are flushed to the backing repository and the remainder of the claim is written there. The limit covers content byte arrays, + * including unused capacity, but not object metadata or buffers owned by processors and their streams. + * + *

+ * Claimant counts for in-memory claims are tracked in the {@link ResourceClaimManager} provided at initialization, the same manager used by the backing + * repository. Each claim that spills holds a single claimant count on its backing claim that this repository owns. Preparing a backing claim retains that ownership + * until {@link #completeBackingClaimTransfer(ContentClaim)} confirms that the NiFi FlowFile Repository references it. If the transfer does not + * complete, purge releases the backing claim for cleanup. In-memory claims that report {@link ResourceClaim#isInUse()} as {@code true} are never handed to the NiFi + * FlowFile Repository for destruction; they are reclaimed by garbage collection once no FlowFile references them, and their memory accounting is released when + * the claim is removed or the repository is purged. + */ +public class SpillableContentRepository implements ContentRepository { + private static final Logger logger = LoggerFactory.getLogger(SpillableContentRepository.class); + + private final ContentRepository backingRepository; + private final FlowFileRepository nifiFlowFileRepository; + private final long memoryThresholdBytes; + private final AtomicLong memoryUsed = new AtomicLong(0L); + private final Set activeClaims = ConcurrentHashMap.newKeySet(); + private final Set backingClaimsPendingCleanup = ConcurrentHashMap.newKeySet(); + + private volatile ResourceClaimManager resourceClaimManager; + + public SpillableContentRepository(final ContentRepository backingRepository, final FlowFileRepository nifiFlowFileRepository, final long memoryThresholdBytes) { + this.backingRepository = backingRepository; + this.nifiFlowFileRepository = nifiFlowFileRepository; + this.memoryThresholdBytes = memoryThresholdBytes; + } + + @Override + public void initialize(final ContentRepositoryContext context) { + this.resourceClaimManager = context.getResourceClaimManager(); + } + + @Override + public void shutdown() { + purge(); + } + + @Override + public Set getContainerNames() { + return backingRepository.getContainerNames(); + } + + @Override + public long getContainerCapacity(final String containerName) throws IOException { + return backingRepository.getContainerCapacity(containerName); + } + + @Override + public long getContainerUsableSpace(final String containerName) throws IOException { + return backingRepository.getContainerUsableSpace(containerName); + } + + @Override + public String getContainerFileStoreName(final String containerName) { + return backingRepository.getContainerFileStoreName(containerName); + } + + @Override + public ContentClaim create(final boolean lossTolerant) { + final SpillableContentClaim contentClaim = new SpillableContentClaim(resourceClaimManager, lossTolerant); + resourceClaimManager.incrementClaimantCount(contentClaim.getResourceClaim()); + activeClaims.add(contentClaim); + return contentClaim; + } + + @Override + public int incrementClaimaintCount(final ContentClaim claim) { + if (claim == null) { + return 0; + } + + if (claim instanceof SpillableContentClaim) { + return resourceClaimManager.incrementClaimantCount(claim.getResourceClaim()); + } + + return backingRepository.incrementClaimaintCount(claim); + } + + @Override + public int getClaimantCount(final ContentClaim claim) { + if (claim == null) { + return 0; + } + + if (claim instanceof SpillableContentClaim) { + return resourceClaimManager.getClaimantCount(claim.getResourceClaim()); + } + + return backingRepository.getClaimantCount(claim); + } + + @Override + public int decrementClaimantCount(final ContentClaim claim) { + if (claim == null) { + return 0; + } + + if (claim instanceof SpillableContentClaim) { + return resourceClaimManager.decrementClaimantCount(claim.getResourceClaim()); + } + + return backingRepository.decrementClaimantCount(claim); + } + + @Override + public boolean remove(final ContentClaim claim) { + if (claim == null) { + return true; + } + + if (!(claim instanceof final SpillableContentClaim spillableClaim)) { + return backingRepository.remove(claim); + } + + releaseClaimForCleanup(spillableClaim); + return true; + } + + @Override + public ContentClaim clone(final ContentClaim original, final boolean lossTolerant) throws IOException { + final ContentClaim clone = create(lossTolerant); + try (final InputStream in = read(original); + final OutputStream out = write(clone)) { + in.transferTo(out); + } catch (final IOException | RuntimeException e) { + decrementClaimantCount(clone); + remove(clone); + throw e; + } + + return clone; + } + + @Override + public long importFrom(final Path content, final ContentClaim claim) throws IOException { + try (final InputStream in = Files.newInputStream(content, StandardOpenOption.READ)) { + return importFrom(in, claim); + } + } + + @Override + public long importFrom(final InputStream content, final ContentClaim claim) throws IOException { + try (final OutputStream out = write(claim)) { + return content.transferTo(out); + } + } + + @Override + public long exportTo(final ContentClaim claim, final Path destination, final boolean append) throws IOException { + final OpenOption[] openOptions = append ? new StandardOpenOption[] {StandardOpenOption.CREATE, StandardOpenOption.APPEND} : + new StandardOpenOption[] {StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING}; + + try (final OutputStream out = Files.newOutputStream(destination, openOptions)) { + return exportTo(claim, out); + } + } + + @Override + public long exportTo(final ContentClaim claim, final Path destination, final boolean append, final long offset, final long length) throws IOException { + final OpenOption[] openOptions = append ? new StandardOpenOption[] {StandardOpenOption.CREATE, StandardOpenOption.APPEND} : + new StandardOpenOption[] {StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING}; + + try (final OutputStream out = Files.newOutputStream(destination, openOptions)) { + return exportTo(claim, out, offset, length); + } + } + + @Override + public long exportTo(final ContentClaim claim, final OutputStream destination) throws IOException { + try (final InputStream in = read(claim)) { + return in.transferTo(destination); + } + } + + @Override + public long exportTo(final ContentClaim claim, final OutputStream destination, final long offset, final long length) throws IOException { + try (final InputStream in = read(claim)) { + StreamUtils.skip(in, offset); + StreamUtils.copy(in, destination, length); + } + + return length; + } + + @Override + public long size(final ContentClaim claim) throws IOException { + if (claim == null) { + return 0; + } + + if (!(claim instanceof final SpillableContentClaim spillableClaim)) { + return backingRepository.size(claim); + } + + final ContentClaim backingClaim = spillableClaim.getBackingClaim(); + if (backingClaim != null) { + return backingRepository.size(backingClaim); + } + + return spillableClaim.getLength(); + } + + @Override + public long size(final ResourceClaim claim) throws IOException { + if (claim instanceof final SpillableResourceClaim spillableResourceClaim) { + final ContentClaim backingClaim = spillableResourceClaim.getBackingClaim(); + return backingClaim == null ? spillableResourceClaim.getLength() : backingRepository.size(backingClaim); + } + + return backingRepository.size(claim); + } + + @Override + public InputStream read(final ContentClaim claim) throws IOException { + if (claim == null) { + return InputStream.nullInputStream(); + } + + if (!(claim instanceof final SpillableContentClaim spillableClaim)) { + return backingRepository.read(claim); + } + + final ContentClaim backingClaim = spillableClaim.getBackingClaim(); + if (backingClaim != null) { + return backingRepository.read(backingClaim); + } + + return spillableClaim.readInMemory(); + } + + @Override + public InputStream read(final ResourceClaim claim) throws IOException { + if (claim instanceof final SpillableResourceClaim spillableResourceClaim) { + final ContentClaim backingClaim = spillableResourceClaim.getBackingClaim(); + return backingClaim == null ? spillableResourceClaim.readInMemory() : backingRepository.read(backingClaim); + } + + return backingRepository.read(claim); + } + + @Override + public OutputStream write(final ContentClaim claim) throws IOException { + if (!(claim instanceof final SpillableContentClaim spillableClaim)) { + return backingRepository.write(claim); + } + + if (!spillableClaim.beginWrite()) { + throw new IllegalStateException("Cannot write to Content Claim because it has already been written to or has an active writer"); + } + + return new SpillableOutputStream(spillableClaim); + } + + @Override + public void purge() { + for (final SpillableContentClaim contentClaim : activeClaims) { + final ResourceClaim resourceClaim = contentClaim.getResourceClaim(); + synchronized (resourceClaim) { + if (resourceClaimManager.getClaimantCount(resourceClaim) == 0 && !contentClaim.isWriteInProgress() && activeClaims.remove(contentClaim)) { + releaseClaimContentsForCleanup(contentClaim); + } + } + } + + submitBackingClaimsForCleanup(); + } + + @Override + public void cleanup() { + } + + @Override + public boolean isAccessible(final ContentClaim claim) throws IOException { + if (claim == null) { + return false; + } + + if (!(claim instanceof final SpillableContentClaim spillableClaim)) { + return backingRepository.isAccessible(claim); + } + + final ContentClaim backingClaim = spillableClaim.getBackingClaim(); + if (backingClaim != null) { + return backingRepository.isAccessible(backingClaim); + } + + return false; + } + + /** + * Returns a claim from the backing Content Repository. For a claim whose content was buffered in memory, the content is written to the backing repository + * and a new backing claim is returned. This repository retains ownership of the backing claim until + * {@link #completeBackingClaimTransfer(ContentClaim)} is called. For any claim that does not belong to this repository, the claim is returned unchanged. + * + * @param claim the claim to prepare + * @return a Content Claim whose content is stored in the backing Content Repository + * @throws IOException if the content cannot be written to the backing repository + */ + ContentClaim prepareBackingClaim(final ContentClaim claim) throws IOException { + if (!(claim instanceof final SpillableContentClaim spillableClaim)) { + // Input FlowFiles can already reference claims from the backing Content Repository. + return claim; + } + + final SpillableResourceClaim resourceClaim = spillableClaim.getResourceClaim(); + synchronized (resourceClaim) { + final ContentClaim existingBackingClaim = spillableClaim.getBackingClaim(); + if (existingBackingClaim != null) { + return existingBackingClaim; + } + + if (spillableClaim.isWriteInProgress()) { + throw new IllegalStateException("Cannot export Content Claim while it is being written"); + } + + final ContentClaim backingClaim = backingRepository.create(spillableClaim.isLossTolerant()); + try (final OutputStream out = backingRepository.write(backingClaim)) { + spillableClaim.writeInMemoryTo(out); + } catch (final IOException | RuntimeException e) { + releaseBackingClaimForCleanup(backingClaim); + throw e; + } + + final long freed = spillableClaim.markExportPrepared(backingClaim); + if (freed > 0) { + memoryUsed.addAndGet(-freed); + } + + return backingClaim; + } + } + + /** + * Hands ownership of a prepared backing claim to the NiFi FlowFile Repository after its records have been updated successfully. + * + * @param claim the original claim passed to {@link #prepareBackingClaim(ContentClaim)} + */ + void completeBackingClaimTransfer(final ContentClaim claim) { + if (!(claim instanceof final SpillableContentClaim spillableClaim)) { + throw new IllegalArgumentException("Content Claim was not created by this repository"); + } + + final SpillableResourceClaim resourceClaim = spillableClaim.getResourceClaim(); + synchronized (resourceClaim) { + final ContentClaim backingClaim = spillableClaim.getBackingClaim(); + if (backingClaim == null) { + return; + } + + if (spillableClaim.transferBackingClaimOwnership()) { + backingRepository.decrementClaimantCount(backingClaim); + } + + activeClaims.remove(spillableClaim); + } + } + + private void releaseClaimForCleanup(final SpillableContentClaim contentClaim) { + final ResourceClaim resourceClaim = contentClaim.getResourceClaim(); + synchronized (resourceClaim) { + if (activeClaims.remove(contentClaim)) { + releaseClaimContentsForCleanup(contentClaim); + } + } + } + + private void releaseClaimContentsForCleanup(final SpillableContentClaim contentClaim) { + final ReleasedContent releasedContent = contentClaim.releaseForCleanup(); + if (releasedContent.memoryBytes() > 0) { + memoryUsed.addAndGet(-releasedContent.memoryBytes()); + } + + if (releasedContent.backingClaim() != null) { + releaseBackingClaimForCleanup(releasedContent.backingClaim()); + } + } + + private void releaseBackingClaimForCleanup(final ContentClaim backingClaim) { + backingRepository.decrementClaimantCount(backingClaim); + backingClaimsPendingCleanup.add(backingClaim); + } + + private synchronized void submitBackingClaimsForCleanup() { + if (backingClaimsPendingCleanup.isEmpty()) { + return; + } + + final Set claimsToDestroy = new HashSet<>(backingClaimsPendingCleanup); + try { + nifiFlowFileRepository.updateRepository(List.of(new StandardRepositoryRecord(claimsToDestroy))); + backingClaimsPendingCleanup.removeAll(claimsToDestroy); + } catch (final IOException e) { + logger.warn("Failed to submit spilled Content Claims for cleanup", e); + } + } + + long getInMemoryByteCount() { + return memoryUsed.get(); + } + + boolean isHeldInMemory(final ContentClaim claim) { + return claim instanceof final SpillableContentClaim spillableClaim && spillableClaim.getBackingClaim() == null && spillableClaim.hasInMemoryContents(); + } + + boolean isSpilled(final ContentClaim claim) { + return claim instanceof final SpillableContentClaim spillableClaim && spillableClaim.getBackingClaim() != null; + } + + private record ReleasedContent(long memoryBytes, ContentClaim backingClaim) { + } + + private final class SpillableOutputStream extends OutputStream { + private final SpillableContentClaim contentClaim; + private final boolean lossTolerant; + private MemoryContents buffer = new MemoryContents(); + private OutputStream spillStream; + private boolean closed = false; + + private SpillableOutputStream(final SpillableContentClaim contentClaim) { + this.contentClaim = contentClaim; + this.lossTolerant = contentClaim.isLossTolerant(); + } + + @Override + public void write(final int value) throws IOException { + if (closed) { + throw new IOException("Cannot write to closed stream"); + } + + if (spillStream == null && !reserveCapacity(1)) { + spillOver(); + } + + if (spillStream == null) { + buffer.write(value); + } else { + spillStream.write(value); + } + } + + @Override + public void write(final byte[] b, final int off, final int len) throws IOException { + Objects.requireNonNull(b); + Objects.checkFromIndexSize(off, len, b.length); + if (closed) { + throw new IOException("Cannot write to closed stream"); + } + + if (len == 0) { + return; + } + + if (spillStream != null) { + spillStream.write(b, off, len); + return; + } + + if (reserveCapacity(len)) { + buffer.write(b, off, len); + return; + } + + spillOver(); + spillStream.write(b, off, len); + } + + private boolean reserveCapacity(final int bytes) { + final long required = bytes - (buffer.capacity() - buffer.size()); + if (required <= 0) { + return true; + } + + final long preferred = Math.max(buffer.nextBufferSize(), required); + while (true) { + final long currentMemoryUsed = memoryUsed.get(); + final long available = memoryThresholdBytes - currentMemoryUsed; + if (available < required) { + return false; + } + + final long reserved = Math.min(preferred, available); + if (memoryUsed.compareAndSet(currentMemoryUsed, currentMemoryUsed + reserved)) { + try { + buffer.allocate(reserved); + } catch (final RuntimeException | Error e) { + memoryUsed.addAndGet(-reserved); + throw e; + } + + return true; + } + } + } + + private void spillOver() throws IOException { + final ContentClaim createdSpillClaim = backingRepository.create(lossTolerant); + OutputStream createdSpillStream = null; + try { + createdSpillStream = backingRepository.write(createdSpillClaim); + buffer.writeTo(createdSpillStream); + contentClaim.markSpilled(createdSpillClaim); + } catch (final IOException | RuntimeException e) { + if (createdSpillStream != null) { + try { + createdSpillStream.close(); + } catch (final Exception closeException) { + e.addSuppressed(closeException); + } + } + + releaseBackingClaimForCleanup(createdSpillClaim); + throw e; + } + + spillStream = createdSpillStream; + final long reserved = buffer.capacity(); + buffer = null; + memoryUsed.addAndGet(-reserved); + logger.debug("Spilled Content Claim {} to the Content Repository after buffering {} bytes in memory; in-memory budget is {} bytes", contentClaim.getResourceClaim().getId(), reserved, + memoryThresholdBytes); + } + + @Override + public void flush() throws IOException { + if (spillStream != null) { + spillStream.flush(); + } + } + + @Override + public void close() throws IOException { + if (closed) { + return; + } + + closed = true; + + try { + if (spillStream != null) { + spillStream.close(); + return; + } + + contentClaim.storeInMemory(buffer); + buffer = null; + } finally { + contentClaim.finishWrite(); + } + } + } + + static final class MemoryContents { + private static final int MAX_BUFFER_SIZE = 64 * 1024; + + private final List buffers = new ArrayList<>(); + private long capacity; + private long size; + private int writeBufferIndex; + private int writeBufferOffset; + + int nextBufferSize() { + return buffers.isEmpty() ? 32 : Math.min(MAX_BUFFER_SIZE, buffers.getLast().length * 2); + } + + void allocate(final long bytes) { + final int originalBufferCount = buffers.size(); + try { + long remaining = bytes; + while (remaining > 0) { + final int bufferSize = (int) Math.min(MAX_BUFFER_SIZE, remaining); + buffers.add(new byte[bufferSize]); + remaining -= bufferSize; + } + } catch (final RuntimeException | Error e) { + buffers.subList(originalBufferCount, buffers.size()).clear(); + throw e; + } + + capacity += bytes; + } + + void write(final int value) { + final byte[] currentBuffer = writableBuffer(); + currentBuffer[writeBufferOffset++] = (byte) value; + size++; + } + + private byte[] writableBuffer() { + if (writeBufferOffset == buffers.get(writeBufferIndex).length) { + writeBufferIndex++; + writeBufferOffset = 0; + } + + return buffers.get(writeBufferIndex); + } + + void write(final byte[] source, final int offset, final int length) { + int remaining = length; + while (remaining > 0) { + final byte[] currentBuffer = writableBuffer(); + final int copied = Math.min(remaining, currentBuffer.length - writeBufferOffset); + System.arraycopy(source, offset + length - remaining, currentBuffer, writeBufferOffset, copied); + writeBufferOffset += copied; + size += copied; + remaining -= copied; + } + } + + long size() { + return size; + } + + long capacity() { + return capacity; + } + + void writeTo(final OutputStream destination) throws IOException { + long remaining = size; + for (final byte[] currentBuffer : buffers) { + final int length = (int) Math.min(remaining, currentBuffer.length); + destination.write(currentBuffer, 0, length); + remaining -= length; + if (remaining == 0) { + break; + } + } + } + + InputStream toInputStream() { + if (size == 0) { + return InputStream.nullInputStream(); + } + + if (buffers.size() == 1) { + return new ByteArrayInputStream(buffers.getFirst(), 0, (int) size); + } + + final Iterator iterator = buffers.iterator(); + final Enumeration streams = new Enumeration<>() { + private long remaining = size; + + @Override + public boolean hasMoreElements() { + return remaining > 0 && iterator.hasNext(); + } + + @Override + public InputStream nextElement() { + final byte[] currentBuffer = iterator.next(); + final int length = (int) Math.min(remaining, currentBuffer.length); + remaining -= length; + return new ByteArrayInputStream(currentBuffer, 0, length); + } + }; + return new SequenceInputStream(streams); + } + } + + private static final class SpillableContentClaim implements ContentClaim { + private final SpillableResourceClaim resourceClaim; + + private SpillableContentClaim(final ResourceClaimManager claimManager, final boolean lossTolerant) { + this.resourceClaim = new SpillableResourceClaim(claimManager, lossTolerant); + } + + @Override + public SpillableResourceClaim getResourceClaim() { + return resourceClaim; + } + + @Override + public long getOffset() { + return 0L; + } + + @Override + public long getLength() { + return resourceClaim.getLength(); + } + + @Override + public boolean isTruncationCandidate() { + return false; + } + + private boolean isLossTolerant() { + return resourceClaim.isLossTolerant(); + } + + private boolean beginWrite() { + return resourceClaim.beginWrite(); + } + + private void finishWrite() { + resourceClaim.finishWrite(); + } + + private boolean isWriteInProgress() { + return resourceClaim.isWriteInProgress(); + } + + private void storeInMemory(final MemoryContents contents) { + resourceClaim.storeInMemory(contents); + } + + private void markSpilled(final ContentClaim backingClaim) { + resourceClaim.markSpilled(backingClaim); + } + + private ContentClaim getBackingClaim() { + return resourceClaim.getBackingClaim(); + } + + private boolean transferBackingClaimOwnership() { + return resourceClaim.transferBackingClaimOwnership(); + } + + private long markExportPrepared(final ContentClaim backingClaim) { + return resourceClaim.markExportPrepared(backingClaim); + } + + private boolean hasInMemoryContents() { + return resourceClaim.hasInMemoryContents(); + } + + private void writeInMemoryTo(final OutputStream out) throws IOException { + resourceClaim.writeInMemoryTo(out); + } + + private InputStream readInMemory() { + return resourceClaim.readInMemory(); + } + + private ReleasedContent releaseForCleanup() { + return resourceClaim.releaseForCleanup(); + } + + @Override + public int compareTo(final ContentClaim o) { + return resourceClaim.compareTo(o.getResourceClaim()); + } + + @Override + public int hashCode() { + return resourceClaim.hashCode(); + } + + @Override + public boolean equals(final Object obj) { + return this == obj; + } + } + + private static final class SpillableResourceClaim extends StandardResourceClaim { + private static final AtomicLong idCounter = new AtomicLong(0L); + + private final ResourceClaimManager claimManager; + private volatile MemoryContents contents; + private volatile ContentClaim backingClaim; + private volatile BackingClaimOwnership backingClaimOwnership = BackingClaimOwnership.NONE; + private volatile boolean writeStarted; + private volatile boolean writeFinished; + private volatile boolean discarded = false; + + private SpillableResourceClaim(final ResourceClaimManager claimManager, final boolean lossTolerant) { + super(claimManager, "in-memory", "in-memory", String.valueOf(idCounter.getAndIncrement()), lossTolerant); + this.claimManager = claimManager; + } + + @Override + public boolean isInUse() { + return true; + } + + private long getLength() { + final ContentClaim currentBackingClaim = backingClaim; + if (currentBackingClaim != null) { + return currentBackingClaim.getLength(); + } + + final MemoryContents currentContents = contents; + return currentContents == null ? 0L : currentContents.size(); + } + + private synchronized boolean beginWrite() { + if (writeStarted || discarded || !isWritable() || contents != null || backingClaim != null) { + return false; + } + + writeStarted = true; + return true; + } + + private synchronized void finishWrite() { + writeFinished = true; + claimManager.freeze(this); + } + + private boolean isWriteInProgress() { + return writeStarted && !writeFinished; + } + + private synchronized void storeInMemory(final MemoryContents contents) { + this.contents = contents; + } + + private synchronized void markSpilled(final ContentClaim backingClaim) { + this.backingClaim = backingClaim; + backingClaimOwnership = BackingClaimOwnership.REPOSITORY; + } + + private ContentClaim getBackingClaim() { + return backingClaim; + } + + private synchronized boolean transferBackingClaimOwnership() { + if (backingClaimOwnership != BackingClaimOwnership.REPOSITORY) { + return false; + } + + backingClaimOwnership = BackingClaimOwnership.EXTERNAL; + return true; + } + + private synchronized long markExportPrepared(final ContentClaim backingClaim) { + this.backingClaim = backingClaim; + backingClaimOwnership = BackingClaimOwnership.REPOSITORY; + + final MemoryContents currentContents = contents; + contents = null; + discarded = true; + claimManager.freeze(this); + return currentContents == null ? 0L : currentContents.capacity(); + } + + private boolean hasInMemoryContents() { + return contents != null; + } + + private void writeInMemoryTo(final OutputStream out) throws IOException { + final MemoryContents currentContents = contents; + if (currentContents != null) { + currentContents.writeTo(out); + } + } + + private InputStream readInMemory() { + final MemoryContents currentContents = contents; + if (currentContents == null) { + return InputStream.nullInputStream(); + } + + return currentContents.toInputStream(); + } + + private synchronized ReleasedContent releaseForCleanup() { + final MemoryContents currentContents = contents; + contents = null; + discarded = true; + claimManager.freeze(this); + + final ContentClaim cleanupClaim; + if (backingClaimOwnership == BackingClaimOwnership.REPOSITORY) { + cleanupClaim = backingClaim; + backingClaimOwnership = BackingClaimOwnership.NONE; + } else { + cleanupClaim = null; + } + + return new ReleasedContent(currentContents == null ? 0L : currentContents.capacity(), cleanupClaim); + } + } + + private enum BackingClaimOwnership { + NONE, + REPOSITORY, + EXTERNAL + } +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/tasks/StatelessFlowTask.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/tasks/StatelessFlowTask.java index a0a60bcdd814..b133bf74c026 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/tasks/StatelessFlowTask.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/tasks/StatelessFlowTask.java @@ -35,6 +35,8 @@ import org.apache.nifi.controller.repository.RepositoryRecordType; import org.apache.nifi.controller.repository.StandardFlowFileRecord; import org.apache.nifi.controller.repository.StandardRepositoryRecord; +import org.apache.nifi.controller.repository.TransferableContentClaimRepository; +import org.apache.nifi.controller.repository.claim.ContentClaim; import org.apache.nifi.controller.repository.metrics.ProcessSessionEventBuilder; import org.apache.nifi.flowfile.FlowFile; import org.apache.nifi.groups.ProcessGroup; @@ -94,6 +96,8 @@ public class StatelessFlowTask { private List cloneResults; private List outputRepositoryRecords; private List cloneProvenanceEvents; + private Set preparedContentClaims; + private List incrementedOutputClaims; // State that is updated during invocation but do not need to be guarded by synchronized block private volatile long shutdownInitiationTime = 0L; @@ -300,7 +304,8 @@ private TriggerResult triggerFlow(final DataflowTriggerContext triggerContext) { } } - private void completeInvocations(final List invocations, final ProvenanceEventRepository statelessProvRepo) throws IOException { + // Visible for testing + void completeInvocations(final List invocations, final ProvenanceEventRepository statelessProvRepo) throws IOException { logger.debug("Completing transactions from {} invocations", invocations.size()); if (invocations.isEmpty()) { return; @@ -329,11 +334,10 @@ private void completeInvocations(final List invocations, final Prove updateClaimantCounts(); // Update the FlowFile Repository to reflect that the input FlowFiles were dropped and the output FlowFiles were created. - try { - updateFlowFileRepository(); - } catch (final Exception e) { - throw new IOException("Failed to update FlowFile Repository after triggering " + this, e); - } + updateFlowFileRepositoryWithRollback(); + + completeContentClaimTransfers(); + incrementedOutputClaims.clear(); updateProvenanceRepository(statelessProvRepo, event -> true); @@ -352,6 +356,8 @@ void resetState() { cloneResults = new ArrayList<>(); outputRepositoryRecords = new ArrayList<>(); cloneProvenanceEvents = new ArrayList<>(); + preparedContentClaims = new HashSet<>(); + incrementedOutputClaims = new ArrayList<>(); } private void failInvocation(final Invocation invocation, final ProvenanceEventRepository statelessProvRepo, final Port destinationPort, final Throwable cause) throws IOException { @@ -412,11 +418,7 @@ private void failInvocation(final Invocation invocation, final ProvenanceEventRe } // Update the FlowFile Repository to reflect that the input FlowFiles were dropped and the output FlowFiles were created. - try { - updateFlowFileRepository(); - } catch (final Exception e) { - throw new IOException("Failed to update FlowFile Repository after triggering " + this, e); - } + updateFlowFileRepositoryWithRollback(); updateProvenanceRepository(statelessProvRepo, event -> eventTypesToKeepOnFailure.contains(event.getEventType())); @@ -469,7 +471,7 @@ public List getCloneProvenanceEvents() { return cloneProvenanceEvents; } - void createOutputRecords(final Map> outputFlowFiles) { + void createOutputRecords(final Map> outputFlowFiles) throws IOException { for (final Map.Entry> entry : outputFlowFiles.entrySet()) { final String portName = entry.getKey(); final Port outputPort = outputPorts.get(portName); @@ -477,19 +479,51 @@ void createOutputRecords(final Map> outputFlowFiles) { final List portFlowFiles = (List) entry.getValue(); final Set outputConnections = outputPort.getConnections(); for (final FlowFileRecord outputFlowFile : portFlowFiles) { - final FlowFileCloneResult cloneResult = ConnectionUtils.clone(outputFlowFile, outputConnections, - nifiFlowFileRepository, null); + // Outbound FlowFiles must reference content that remains available after Stateless content is purged. + final FlowFileRecord outputFlowFileWithBackingClaim = prepareBackingClaim(outputFlowFile); + + final FlowFileCloneResult cloneResult = ConnectionUtils.clone(outputFlowFileWithBackingClaim, outputConnections, nifiFlowFileRepository, null); cloneResults.add(cloneResult); final List repoRecords = cloneResult.getRepositoryRecords(); outputRepositoryRecords.addAll(repoRecords); // If we generated any clones, create provenance events for them. - createCloneProvenanceEvent(outputFlowFile, repoRecords, outputPort).ifPresent(cloneProvenanceEvents::add); + createCloneProvenanceEvent(outputFlowFileWithBackingClaim, repoRecords, outputPort).ifPresent(cloneProvenanceEvents::add); } } } + private FlowFileRecord prepareBackingClaim(final FlowFileRecord flowFile) throws IOException { + final ContentClaim contentClaim = flowFile.getContentClaim(); + if (contentClaim == null || !(nifiContentRepository instanceof final TransferableContentClaimRepository transferableContentClaimRepository)) { + return flowFile; + } + + final ContentClaim backingClaim = transferableContentClaimRepository.prepareBackingClaim(contentClaim); + if (backingClaim == contentClaim) { + return flowFile; + } + + preparedContentClaims.add(contentClaim); + return new StandardFlowFileRecord.Builder() + .fromFlowFile(flowFile) + .contentClaim(backingClaim) + .build(); + } + + private void completeContentClaimTransfers() { + if (!(nifiContentRepository instanceof final TransferableContentClaimRepository transferableContentClaimRepository)) { + return; + } + + for (final ContentClaim preparedContentClaim : preparedContentClaims) { + transferableContentClaimRepository.completeBackingClaimTransfer(preparedContentClaim); + } + + preparedContentClaims.clear(); + } + void updateProvenanceRepository(final ProvenanceEventRepository statelessRepo, final Predicate eventFilter) { long firstProvEventId = 0; @@ -539,11 +573,30 @@ void updateClaimantCounts() { // for the input FlowFile is simpler and will work just as well. for (final RepositoryRecord outputRepoRecord : outputRepositoryRecords) { if (outputRepoRecord.getType() != RepositoryRecordType.DELETE) { - nifiContentRepository.incrementClaimaintCount(outputRepoRecord.getCurrentClaim()); + final ContentClaim contentClaim = outputRepoRecord.getCurrentClaim(); + nifiContentRepository.incrementClaimaintCount(contentClaim); + incrementedOutputClaims.add(contentClaim); } } } + private void rollbackClaimantCounts() { + for (final ContentClaim contentClaim : incrementedOutputClaims) { + nifiContentRepository.decrementClaimantCount(contentClaim); + } + + incrementedOutputClaims.clear(); + } + + private void updateFlowFileRepositoryWithRollback() throws IOException { + try { + updateFlowFileRepository(); + } catch (final Exception e) { + rollbackClaimantCounts(); + throw new IOException("Failed to update FlowFile Repository after triggering " + this, e); + } + } + private void updateFlowFileRepository() throws IOException { // Update the FlowFile repository nifiFlowFileRepository.updateRepository(outputRepositoryRecords); diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/DeferredStatelessContentRepositoryTest.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/DeferredStatelessContentRepositoryTest.java new file mode 100644 index 000000000000..f3ada89b79f6 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/DeferredStatelessContentRepositoryTest.java @@ -0,0 +1,66 @@ +/* + * 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.nifi.controller.repository; + +import org.apache.nifi.controller.repository.claim.ContentClaim; +import org.apache.nifi.controller.repository.claim.ResourceClaim; +import org.apache.nifi.controller.repository.claim.ResourceClaimManager; +import org.apache.nifi.controller.repository.claim.StandardResourceClaimManager; +import org.apache.nifi.events.EventReporter; +import org.apache.nifi.groups.ProcessGroup; +import org.apache.nifi.stateless.repository.ByteArrayContentRepository; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.concurrent.atomic.AtomicLong; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +class DeferredStatelessContentRepositoryTest { + + @Test + void testResolvesChangedThreshold() throws IOException { + final ResourceClaimManager resourceClaimManager = new StandardResourceClaimManager(); + final ContentRepository backingRepository = new ByteArrayContentRepository(); + backingRepository.initialize(new StandardContentRepositoryContext(resourceClaimManager, EventReporter.NO_OP)); + final FlowFileRepository flowFileRepository = mock(FlowFileRepository.class); + final ProcessGroup processGroup = mock(ProcessGroup.class); + final AtomicLong threshold = new AtomicLong(100L); + when(processGroup.resolveStatelessContentMaxHeap()).thenAnswer(invocation -> threshold.get()); + final DeferredStatelessContentRepository repository = new DeferredStatelessContentRepository( + processGroup, backingRepository, flowFileRepository, resourceClaimManager, EventReporter.NO_OP); + + final ContentClaim inMemoryClaim = repository.create(false); + assertEquals("in-memory", inMemoryClaim.getResourceClaim().getContainer()); + repository.decrementClaimantCount(inMemoryClaim); + repository.purge(); + + final ResourceClaim resourceClaim = inMemoryClaim.getResourceClaim(); + assertNull(resourceClaimManager.getResourceClaim(resourceClaim.getContainer(), resourceClaim.getSection(), resourceClaim.getId())); + assertEquals(0, resourceClaimManager.getClaimantCount(resourceClaim)); + assertFalse(resourceClaim.isWritable()); + + threshold.set(0L); + final ContentClaim backingClaim = repository.create(false); + assertEquals("container", backingClaim.getResourceClaim().getContainer()); + } +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/SpillableContentRepositoryTest.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/SpillableContentRepositoryTest.java new file mode 100644 index 000000000000..16fdc06e9955 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/SpillableContentRepositoryTest.java @@ -0,0 +1,557 @@ +/* + * 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.nifi.controller.repository; + +import org.apache.nifi.controller.repository.SpillableContentRepository.MemoryContents; +import org.apache.nifi.controller.repository.claim.ContentClaim; +import org.apache.nifi.controller.repository.claim.ResourceClaim; +import org.apache.nifi.controller.repository.claim.ResourceClaimManager; +import org.apache.nifi.controller.repository.claim.StandardResourceClaimManager; +import org.apache.nifi.events.EventReporter; +import org.apache.nifi.stateless.repository.ByteArrayContentRepository; +import org.apache.nifi.stream.io.ByteCountingOutputStream; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.mockito.ArgumentCaptor; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.FilterOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.lang.reflect.Field; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +class SpillableContentRepositoryTest { + private static final long BUDGET = 100L; + + private ResourceClaimManager resourceClaimManager; + private TrackingContentRepository backingRepository; + private FlowFileRepository flowFileRepository; + private SpillableContentRepository repository; + + @BeforeEach + void setup() throws IOException { + resourceClaimManager = new StandardResourceClaimManager(); + backingRepository = new TrackingContentRepository(); + backingRepository.initialize(new StandardContentRepositoryContext(resourceClaimManager, EventReporter.NO_OP)); + flowFileRepository = mock(FlowFileRepository.class); + repository = createRepository(resourceClaimManager, BUDGET); + } + + @ParameterizedTest + @ValueSource(ints = {1, 200_000}) + void testInMemoryContentSupportsRepeatedReads(final int length) throws IOException { + repository = createRepository(resourceClaimManager, Math.max(BUDGET, length)); + final byte[] content = bytes(length); + final ContentClaim claim = repository.create(false); + write(claim, content); + + assertTrue(repository.isHeldInMemory(claim)); + assertFalse(repository.isSpilled(claim)); + assertFalse(repository.isAccessible(claim)); + assertEquals(length, repository.size(claim)); + assertEquals(Math.max(32, length), repository.getInMemoryByteCount()); + assertArrayEquals(content, readAll(claim)); + assertArrayEquals(content, readAll(claim)); + final ContentClaim backingClaim = repository.prepareBackingClaim(claim); + assertArrayEquals(content, readAll(backingClaim)); + assertEquals(0L, repository.getInMemoryByteCount()); + } + + @Test + void testSmallClaimsAccountForAllocatedCapacity() throws IOException { + final byte[] content = new byte[] {1}; + for (int claimIndex = 0; claimIndex < 5; claimIndex++) { + final ContentClaim claim = repository.create(false); + write(claim, content); + + assertArrayEquals(content, readAll(claim)); + assertEquals(claimIndex < 4, repository.isHeldInMemory(claim)); + assertEquals(claimIndex == 4, repository.isSpilled(claim)); + assertEquals(Math.min((claimIndex + 1) * 32L, BUDGET), repository.getInMemoryByteCount()); + } + } + + @Test + void testConcurrentWritersShareCapacityBudget() throws Exception { + final List> futures = new ArrayList<>(); + final CountDownLatch startWriting = new CountDownLatch(1); + try (final ExecutorService executor = Executors.newFixedThreadPool(8)) { + for (int claimIndex = 0; claimIndex < 32; claimIndex++) { + futures.add(executor.submit(() -> { + final ContentClaim claim = repository.create(false); + assertTrue(startWriting.await(10, TimeUnit.SECONDS)); + try (final OutputStream output = repository.write(claim)) { + output.write(bytes(32)); + } + + return claim; + })); + } + + startWriting.countDown(); + int inMemoryClaimCount = 0; + for (final Future future : futures) { + final ContentClaim claim = future.get(10, TimeUnit.SECONDS); + assertArrayEquals(bytes(32), readAll(claim)); + if (repository.isHeldInMemory(claim)) { + inMemoryClaimCount++; + } + + repository.decrementClaimantCount(claim); + } + + assertEquals(3, inMemoryClaimCount); + assertEquals(96L, repository.getInMemoryByteCount()); + } + + repository.purge(); + assertEquals(0L, repository.getInMemoryByteCount()); + } + + @Test + void testSegmentedContentBeyondIntegerLength() throws Exception { + final int segmentSize = 64 * 1024; + final int segmentCount = (int) ((1L + Integer.MAX_VALUE) / segmentSize); + final MemoryContents contents = new MemoryContents(); + contents.allocate(2L * segmentSize); + final List buffers = getBuffers(contents); + assertEquals(2, buffers.size()); + final byte[] segment = buffers.getFirst(); + assertEquals(segmentSize, segment.length); + assertEquals(segmentSize, buffers.getLast().length); + Arrays.fill(segment, (byte) 1); + buffers.addAll(Collections.nCopies(segmentCount - 2, segment)); + setField(contents, "capacity", 1L + Integer.MAX_VALUE); + setField(contents, "size", (long) Integer.MAX_VALUE); + setField(contents, "writeBufferIndex", segmentCount - 1); + setField(contents, "writeBufferOffset", segmentSize - 1); + + contents.write(2); + contents.allocate(32); + contents.write(new byte[] {3, 4}, 0, 2); + assertEquals(3L + Integer.MAX_VALUE, contents.size()); + assertEquals(33L + Integer.MAX_VALUE, contents.capacity()); + + try (final InputStream input = contents.toInputStream()) { + input.skipNBytes(Integer.MAX_VALUE); + assertArrayEquals(new byte[] {2, 3, 4}, input.readAllBytes()); + } + + try (final ByteCountingOutputStream exported = new ByteCountingOutputStream(OutputStream.nullOutputStream())) { + contents.writeTo(exported); + assertEquals(contents.size(), exported.getBytesWritten()); + } + } + + @Test + void testMultipleWritesCrossingBudgetSpillMidStream() throws IOException { + final byte[] content = bytes(120); + final ContentClaim claim = repository.create(false); + try (final OutputStream out = repository.write(claim)) { + out.write(content, 0, 32); + out.write(content[32]); + assertEquals(96L, repository.getInMemoryByteCount()); + out.write(content, 33, 27); + out.write(content, 60, content.length - 60); + } + + assertTrue(repository.isSpilled(claim)); + assertEquals(0L, repository.getInMemoryByteCount()); + assertEquals(120L, repository.size(claim)); + + assertArrayEquals(content, readAll(claim)); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void testInMemoryCleanupReleasesCapacityOnce(final boolean purge) throws Exception { + final ContentClaim claim = repository.create(false); + write(claim, bytes(1)); + assertEquals(32L, repository.getInMemoryByteCount()); + + assertEquals(0, repository.decrementClaimantCount(claim)); + if (purge) { + repository.purge(); + } else { + assertTrue(repository.remove(claim)); + } + + assertEquals(0L, repository.getInMemoryByteCount()); + assertClaimNotTracked(claim); + repository.remove(claim); + repository.purge(); + assertEquals(0L, repository.getInMemoryByteCount()); + verify(flowFileRepository, never()).updateRepository(anyList()); + } + + @Test + void testPurgeCleansOnlyUnreferencedClaims() throws Exception { + final byte[] inMemoryContent = bytes(1); + final ContentClaim inMemoryClaim = repository.create(false); + assertEquals(1, repository.getClaimantCount(inMemoryClaim)); + assertEquals(2, repository.incrementClaimaintCount(inMemoryClaim)); + write(inMemoryClaim, inMemoryContent); + assertEquals(2, repository.getClaimantCount(inMemoryClaim)); + assertFalse(inMemoryClaim.getResourceClaim().isWritable()); + assertEquals(1, repository.decrementClaimantCount(inMemoryClaim)); + + final byte[] spilledContent = bytes(150); + final ContentClaim spilledClaim = repository.create(false); + write(spilledClaim, spilledContent); + + repository.purge(); + + verify(flowFileRepository, never()).updateRepository(anyList()); + final ContentClaim discardedMemoryClaim = repository.create(false); + write(discardedMemoryClaim, bytes(1)); + repository.decrementClaimantCount(discardedMemoryClaim); + final ContentClaim discardedSpilledClaim = repository.create(false); + write(discardedSpilledClaim, bytes(160)); + repository.decrementClaimantCount(discardedSpilledClaim); + + repository.purge(); + + assertEquals(32L, repository.getInMemoryByteCount()); + assertArrayEquals(inMemoryContent, readAll(inMemoryClaim)); + assertArrayEquals(spilledContent, readAll(spilledClaim)); + final ArgumentCaptor> recordsCaptor = captureRepositoryRecords(); + verify(flowFileRepository).updateRepository(recordsCaptor.capture()); + assertEquals(1, recordsCaptor.getValue().size()); + final RepositoryRecord record = recordsCaptor.getValue().getFirst(); + assertEquals(RepositoryRecordType.CLEANUP_TRANSIENT_CLAIMS, record.getType()); + assertEquals(1, record.getTransientClaims().size()); + final ContentClaim cleanupClaim = record.getTransientClaims().getFirst(); + assertArrayEquals(bytes(160), readAll(cleanupClaim)); + assertEquals(0, backingRepository.getClaimantCount(cleanupClaim)); + assertClaimNotTracked(discardedMemoryClaim); + assertClaimNotTracked(discardedSpilledClaim); + assertEquals(0, repository.decrementClaimantCount(inMemoryClaim)); + assertClaimNotTracked(inMemoryClaim); + } + + @Test + void testPurgeWhileClaimIsBeingCreated() throws Exception { + final ResourceClaimManager blockingResourceClaimManager = spy(new StandardResourceClaimManager()); + final CountDownLatch incrementStarted = new CountDownLatch(1); + final CountDownLatch allowIncrement = new CountDownLatch(1); + doAnswer(invocation -> { + incrementStarted.countDown(); + if (!allowIncrement.await(10, TimeUnit.SECONDS)) { + throw new IllegalStateException("Timed out waiting to increment claimant count"); + } + + return invocation.callRealMethod(); + }).when(blockingResourceClaimManager).incrementClaimantCount(any(ResourceClaim.class)); + + backingRepository.initialize(new StandardContentRepositoryContext(blockingResourceClaimManager, EventReporter.NO_OP)); + final SpillableContentRepository concurrentRepository = createRepository(blockingResourceClaimManager, BUDGET); + + final ExecutorService executorService = Executors.newSingleThreadExecutor(); + try { + final Future claimFuture = executorService.submit(() -> concurrentRepository.create(false)); + assertTrue(incrementStarted.await(10, TimeUnit.SECONDS)); + concurrentRepository.purge(); + allowIncrement.countDown(); + + final ContentClaim claim = claimFuture.get(10, TimeUnit.SECONDS); + try (final OutputStream out = concurrentRepository.write(claim)) { + out.write(bytes(10)); + } + + assertArrayEquals(bytes(10), readAll(concurrentRepository, claim)); + } finally { + allowIncrement.countDown(); + executorService.shutdownNow(); + } + } + + @Test + void testPurgePreservesClaimWithOpenWriter() throws Exception { + final byte[] content = bytes(50); + final ContentClaim claim = repository.create(false); + try (final OutputStream out = repository.write(claim)) { + assertThrows(IllegalStateException.class, () -> repository.write(claim)); + out.write(content); + repository.decrementClaimantCount(claim); + repository.purge(); + assertEquals(content.length, repository.getInMemoryByteCount()); + } + + assertArrayEquals(content, readAll(claim)); + assertThrows(IllegalStateException.class, () -> repository.write(claim)); + assertClaimNotTracked(claim); + assertFalse(claim.getResourceClaim().isWritable()); + repository.purge(); + assertEquals(0L, repository.getInMemoryByteCount()); + } + + @ParameterizedTest + @ValueSource(ints = {0, 1, 150}) + void testBackingClaimTransferCompletesOnce(final int length) throws IOException { + final byte[] content = bytes(length); + final ContentClaim claim = repository.create(false); + write(claim, content); + + final boolean spilled = length > BUDGET; + assertEquals(spilled, repository.isSpilled(claim)); + assertEquals(!spilled, repository.isHeldInMemory(claim)); + assertEquals(spilled, repository.isAccessible(claim)); + assertEquals(length, repository.size(claim)); + assertEquals(length, repository.size(claim.getResourceClaim())); + assertArrayEquals(content, readAll(claim)); + try (final InputStream input = repository.read(claim.getResourceClaim())) { + assertArrayEquals(content, input.readAllBytes()); + } + + assertEquals(length == 1 ? 32L : 0L, repository.getInMemoryByteCount()); + + final ContentClaim backingClaim = repository.prepareBackingClaim(claim); + assertNotSame(claim, backingClaim); + assertSame(backingClaim, repository.prepareBackingClaim(claim)); + assertArrayEquals(content, readAll(backingClaim)); + assertEquals(length, backingClaim.getLength()); + assertEquals(0L, repository.getInMemoryByteCount()); + assertEquals(1, backingRepository.writeCount.get()); + assertEquals(1, backingRepository.closedStreamCount.get()); + assertEquals(1, repository.getClaimantCount(backingClaim)); + + repository.incrementClaimaintCount(backingClaim); + repository.completeBackingClaimTransfer(claim); + repository.completeBackingClaimTransfer(claim); + assertEquals(1, repository.getClaimantCount(backingClaim)); + assertThrows(IllegalArgumentException.class, () -> repository.completeBackingClaimTransfer(backingClaim)); + repository.decrementClaimantCount(claim); + repository.purge(); + assertArrayEquals(content, readAll(backingClaim)); + verify(flowFileRepository, never()).updateRepository(anyList()); + } + + @Test + void testPreparingBackingClaimFailureSubmitsClaimForCleanup() throws IOException { + final ContentClaim claim = repository.create(false); + write(claim, bytes(50)); + backingRepository.failNextWrite = true; + + assertThrows(IOException.class, () -> repository.prepareBackingClaim(claim)); + + repository.purge(); + verify(flowFileRepository).updateRepository(anyList()); + } + + @Test + void testPurgeCleansBackingClaimWhenRepositoryUpdateFails() throws IOException { + final ContentClaim claim = repository.create(false); + write(claim, bytes(50)); + + final ContentClaim backingClaim = repository.prepareBackingClaim(claim); + repository.incrementClaimaintCount(backingClaim); + repository.decrementClaimantCount(backingClaim); + repository.decrementClaimantCount(claim); + repository.purge(); + + final ArgumentCaptor> recordsCaptor = captureRepositoryRecords(); + verify(flowFileRepository).updateRepository(recordsCaptor.capture()); + assertEquals(List.of(backingClaim), recordsCaptor.getValue().getFirst().getTransientClaims()); + assertEquals(0, resourceClaimManager.getClaimantCount(backingClaim.getResourceClaim())); + } + + @Test + void testPurgeRetriesFailedCleanupSubmission() throws IOException { + final ContentClaim claim = repository.create(false); + write(claim, bytes(150)); + repository.decrementClaimantCount(claim); + doThrow(new IOException("First update failed")).doNothing().when(flowFileRepository).updateRepository(anyList()); + + repository.purge(); + repository.purge(); + + verify(flowFileRepository, times(2)).updateRepository(anyList()); + } + + @Test + void testSpillFailureRetainsBufferedContentAndMemoryAccounting() throws IOException { + final byte[] content = bytes(60); + final ContentClaim claim = repository.create(false); + final OutputStream out = repository.write(claim); + out.write(content); + backingRepository.failNextWrite = true; + + assertThrows(IOException.class, () -> out.write(bytes(60))); + out.close(); + + assertEquals(60L, repository.getInMemoryByteCount()); + assertArrayEquals(content, readAll(claim)); + } + + @Test + void testContentOperationsAcrossStorageStates(@TempDir final Path tempDirectory) throws IOException { + for (final int length : new int[] {40, 150}) { + final byte[] content = bytes(length); + final ContentClaim claim = repository.create(false); + assertEquals(length, repository.importFrom(new ByteArrayInputStream(content), claim)); + + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + assertEquals(length, repository.exportTo(claim, out)); + assertArrayEquals(content, out.toByteArray()); + + final Path destination = tempDirectory.resolve("content-" + length); + Files.write(destination, bytes(length * 2)); + assertEquals(length, repository.exportTo(claim, destination, false)); + assertArrayEquals(content, Files.readAllBytes(destination)); + + final ContentClaim clone = repository.clone(claim, false); + assertArrayEquals(content, readAll(clone)); + assertEquals(length > BUDGET, repository.isSpilled(clone)); + assertEquals(length <= BUDGET, repository.isHeldInMemory(clone)); + } + } + + private SpillableContentRepository createRepository(final ResourceClaimManager claimManager, final long memoryThresholdBytes) { + final SpillableContentRepository contentRepository = new SpillableContentRepository(backingRepository, flowFileRepository, memoryThresholdBytes); + contentRepository.initialize(new StandardContentRepositoryContext(claimManager, EventReporter.NO_OP)); + return contentRepository; + } + + private ArgumentCaptor> captureRepositoryRecords() { + @SuppressWarnings("unchecked") + final ArgumentCaptor> captor = ArgumentCaptor.forClass(List.class); + return captor; + } + + private void write(final ContentClaim claim, final byte[] content) throws IOException { + try (final OutputStream out = repository.write(claim)) { + out.write(content); + } + } + + private byte[] readAll(final ContentClaim claim) throws IOException { + return readAll(repository, claim); + } + + private byte[] readAll(final ContentRepository contentRepository, final ContentClaim claim) throws IOException { + try (final InputStream in = contentRepository.read(claim)) { + return in.readAllBytes(); + } + } + + private void assertClaimNotTracked(final ContentClaim claim) { + final ResourceClaim resourceClaim = claim.getResourceClaim(); + assertNull(resourceClaimManager.getResourceClaim(resourceClaim.getContainer(), resourceClaim.getSection(), resourceClaim.getId())); + assertEquals(0, resourceClaimManager.getClaimantCount(resourceClaim)); + } + + @SuppressWarnings("unchecked") + private List getBuffers(final MemoryContents contents) throws ReflectiveOperationException { + final Field buffersField = MemoryContents.class.getDeclaredField("buffers"); + buffersField.setAccessible(true); + return (List) buffersField.get(contents); + } + + private void setField(final MemoryContents contents, final String name, final Object value) throws ReflectiveOperationException { + final Field field = MemoryContents.class.getDeclaredField(name); + field.setAccessible(true); + field.set(contents, value); + } + + private static byte[] bytes(final int length) { + final byte[] content = new byte[length]; + for (int i = 0; i < length; i++) { + content[i] = (byte) ('A' + (i % 26)); + } + + return content; + } + + private static final class TrackingContentRepository extends ByteArrayContentRepository { + private final AtomicInteger writeCount = new AtomicInteger(0); + private final AtomicInteger closedStreamCount = new AtomicInteger(0); + private boolean failNextWrite; + + @Override + public OutputStream write(final ContentClaim claim) { + writeCount.incrementAndGet(); + final boolean failWrite = failNextWrite; + failNextWrite = false; + return new FilterOutputStream(super.write(claim)) { + @Override + public void write(final int value) throws IOException { + if (failWrite) { + throw new IOException("Write failed"); + } + + out.write(value); + } + + @Override + public void write(final byte[] bytes, final int offset, final int length) throws IOException { + if (failWrite) { + throw new IOException("Write failed"); + } + + out.write(bytes, offset, length); + } + + @Override + public void close() throws IOException { + super.close(); + closedStreamCount.incrementAndGet(); + } + }; + } + + @Override + public boolean isAccessible(final ContentClaim contentClaim) { + return contentClaim != null && getBytes(contentClaim) != null; + } + } +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/service/mock/MockProcessGroup.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/service/mock/MockProcessGroup.java index 77c6e0391c0a..2dc44faeff3c 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/service/mock/MockProcessGroup.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/service/mock/MockProcessGroup.java @@ -909,6 +909,37 @@ public String getStatelessFlowTimeout() { return null; } + @Override + public String getStatelessContentMaxHeap() { + return null; + } + + @Override + public void setStatelessContentMaxHeap(final String maxSize) { + } + + @Override + public Integer getStatelessContentMaxHeapPercentage() { + return null; + } + + @Override + public void setStatelessContentMaxHeapPercentage(final Integer heapPercentage) { + } + + @Override + public long resolveStatelessContentMaxHeap() { + return 0L; + } + + @Override + public void verifyCanSetStatelessContentMaxHeap(final String maxSize) { + } + + @Override + public void verifyCanSetStatelessContentMaxHeapPercentage(final Integer heapPercentage) { + } + @Override public FlowFileActivity getFlowFileActivity() { return new ProcessGroupFlowFileActivity(this); diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/tasks/TestStatelessFlowTask.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/tasks/TestStatelessFlowTask.java index 203360823415..59d3292c9d4f 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/tasks/TestStatelessFlowTask.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/tasks/TestStatelessFlowTask.java @@ -27,11 +27,13 @@ import org.apache.nifi.controller.metrics.ProcessSessionEvent; import org.apache.nifi.controller.queue.FlowFileQueue; import org.apache.nifi.controller.repository.ContentRepository; +import org.apache.nifi.controller.repository.DeferredStatelessContentRepository; import org.apache.nifi.controller.repository.FlowFileEventRepository; import org.apache.nifi.controller.repository.FlowFileRecord; import org.apache.nifi.controller.repository.FlowFileRepository; import org.apache.nifi.controller.repository.RepositoryRecord; import org.apache.nifi.controller.repository.RepositoryRecordType; +import org.apache.nifi.controller.repository.StandardContentRepositoryContext; import org.apache.nifi.controller.repository.claim.ContentClaim; import org.apache.nifi.controller.repository.claim.ResourceClaim; import org.apache.nifi.controller.repository.claim.ResourceClaimManager; @@ -41,6 +43,7 @@ import org.apache.nifi.controller.service.mock.MockProcessGroup; import org.apache.nifi.controller.tasks.StatelessFlowTask.Invocation; import org.apache.nifi.controller.tasks.StatelessFlowTask.PolledFlowFile; +import org.apache.nifi.events.EventReporter; import org.apache.nifi.flowfile.FlowFile; import org.apache.nifi.flowfile.attributes.CoreAttributes; import org.apache.nifi.groups.ProcessGroup; @@ -53,11 +56,15 @@ import org.apache.nifi.provenance.StandardProvenanceEventRecord; import org.apache.nifi.stateless.flow.StatelessDataflow; import org.apache.nifi.stateless.flow.TriggerResult; +import org.apache.nifi.stateless.repository.ByteArrayContentRepository; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -68,15 +75,19 @@ import java.util.OptionalLong; import java.util.Set; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -94,6 +105,11 @@ public class TestStatelessFlowTask { private ProvenanceEventRepository statelessProvRepo; private FlowFileActivity groupNodeFlowFileActivity; private StatelessDataflow statelessFlow; + private StatelessGroupNode statelessGroupNode; + private FlowFileRepository flowFileRepository; + private FlowFileEventRepository flowFileEventRepository; + private ProvenanceEventRepository provenanceRepository; + private ComponentLog logger; @BeforeEach public void setup() throws IOException { @@ -130,12 +146,12 @@ public void setup() throws IOException { statelessFlow = mock(StatelessDataflow.class); when(statelessFlow.getLatestActivityTime()).thenReturn(OptionalLong.empty()); - final StatelessGroupNode statelessGroupNode = mock(StatelessGroupNode.class); + statelessGroupNode = mock(StatelessGroupNode.class); when(statelessGroupNode.getProcessGroup()).thenReturn(rootGroup); groupNodeFlowFileActivity = new ConnectableFlowFileActivity(); when(statelessGroupNode.getFlowFileActivity()).thenReturn(groupNodeFlowFileActivity); - final FlowFileRepository flowFileRepo = mock(FlowFileRepository.class); + flowFileRepository = mock(FlowFileRepository.class); resourceClaimManager = new StandardResourceClaimManager(); final ContentRepository contentRepo = mock(ContentRepository.class); @@ -158,32 +174,23 @@ public void setup() throws IOException { }).when(contentRepo).decrementClaimantCount(any(ContentClaim.class)); flowFileEventsByComponentId = new HashMap<>(); - final FlowFileEventRepository eventRepository = mock(FlowFileEventRepository.class); + flowFileEventRepository = mock(FlowFileEventRepository.class); doAnswer(invocation -> { final ProcessSessionEvent event = invocation.getArgument(0, ProcessSessionEvent.class); flowFileEventsByComponentId.put(event.getComponentMetricContext().id(), event); return null; - }).when(eventRepository).updateRepository(any(ProcessSessionEvent.class)); - final ComponentLog logger = new MockComponentLogger(); + }).when(flowFileEventRepository).updateRepository(any(ProcessSessionEvent.class)); + logger = new MockComponentLogger(); registeredProvenanceEvents = new ArrayList<>(); - final ProvenanceEventRepository provenanceRepo = mock(ProvenanceEventRepository.class); + provenanceRepository = mock(ProvenanceEventRepository.class); doAnswer(invocation -> { final Iterable events = invocation.getArgument(0, Iterable.class); events.forEach(registeredProvenanceEvents::add); return null; - }).when(provenanceRepo).registerEvents(any(Iterable.class)); - - task = new StatelessFlowTask.Builder() - .statelessFlow(statelessFlow) - .statelessGroupNode(statelessGroupNode) - .nifiFlowFileRepository(flowFileRepo) - .nifiContentRepository(contentRepo) - .nifiProvenanceRepository(provenanceRepo) - .flowFileEventRepository(eventRepository) - .logger(logger) - .build(); + }).when(provenanceRepository).registerEvents(any(Iterable.class)); + task = createTask(contentRepo); task.resetState(); } @@ -216,7 +223,7 @@ public void testDropInputFlowFile() { } @Test - public void testCreateOutputRecordsOnSuccessWithOneOutput() { + public void testCreateOutputRecordsOnSuccessWithOneOutput() throws IOException { final Connection conn1 = mockConnection(); final Set singleConnectionSet = Collections.singleton(conn1); @@ -242,7 +249,89 @@ public void testCreateOutputRecordsOnSuccessWithOneOutput() { } @Test - public void testCreateOutputRecordsOnSuccessWithOnePortTwoConnections() { + public void testCreateOutputRecordsExportsSharedInMemoryClaimOnce() throws IOException { + final byte[] content = "shared content".getBytes(StandardCharsets.UTF_8); + final DeferredStatelessContentRepository deferredRepository = createDeferredContentRepository(); + + final ContentClaim inMemoryClaim = deferredRepository.create(false); + try (final OutputStream out = deferredRepository.write(inMemoryClaim)) { + out.write(content); + } + + final int firstFlowFileOffset = 5; + final FlowFileRecord firstFlowFile = new MockFlowFileRecord(Map.of(), content.length - firstFlowFileOffset, inMemoryClaim) { + @Override + public long getContentClaimOffset() { + return firstFlowFileOffset; + } + }; + final FlowFileRecord secondFlowFile = new MockFlowFileRecord(Map.of(), content.length, inMemoryClaim); + deferredRepository.incrementClaimaintCount(inMemoryClaim); + + final Connection connection = mockConnection(); + when(successPort.getConnections()).thenReturn(Set.of(connection)); + final StatelessFlowTask boundaryTask = createTask(deferredRepository); + boundaryTask.resetState(); + boundaryTask.createOutputRecords(Map.of("success", List.of(firstFlowFile, secondFlowFile))); + + final List outputRecords = boundaryTask.getOutputRepositoryRecords(); + final ContentClaim firstBackingClaim = outputRecords.get(0).getCurrentClaim(); + final ContentClaim secondBackingClaim = outputRecords.get(1).getCurrentClaim(); + assertSame(firstBackingClaim, secondBackingClaim); + final FlowFileRecord firstOutputFlowFile = outputRecords.get(0).getCurrent(); + assertEquals(firstFlowFileOffset, firstOutputFlowFile.getContentClaimOffset()); + try (final InputStream in = deferredRepository.read(firstBackingClaim)) { + in.skipNBytes(firstOutputFlowFile.getContentClaimOffset()); + assertArrayEquals(Arrays.copyOfRange(content, firstFlowFileOffset, content.length), in.readNBytes((int) firstOutputFlowFile.getSize())); + } + try (final InputStream in = deferredRepository.read(secondBackingClaim)) { + assertArrayEquals(content, in.readAllBytes()); + } + } + + @Test + public void testCompleteInvocationsRollsBackPreparedExportsWhenRepositoryUpdateFails() throws IOException { + final DeferredStatelessContentRepository deferredRepository = createDeferredContentRepository(); + + final ContentClaim firstClaim = deferredRepository.create(false); + final ContentClaim secondClaim = deferredRepository.create(false); + try (final OutputStream out = deferredRepository.write(firstClaim)) { + out.write("first".getBytes(StandardCharsets.UTF_8)); + } + try (final OutputStream out = deferredRepository.write(secondClaim)) { + out.write("second".getBytes(StandardCharsets.UTF_8)); + } + + final FlowFileRecord firstFlowFile = new MockFlowFileRecord(Map.of(), 5, firstClaim); + final FlowFileRecord secondFlowFile = new MockFlowFileRecord(Map.of(), 6, secondClaim); + final TriggerResult triggerResult = mock(TriggerResult.class); + when(triggerResult.getOutputFlowFiles()).thenReturn(Map.of("success", List.of(firstFlowFile, secondFlowFile))); + final Invocation invocation = new Invocation(); + invocation.setTriggerResult(triggerResult); + + final Connection connection = mockConnection(); + when(successPort.getConnections()).thenReturn(Set.of(connection)); + doThrow(new IOException("Repository update failed")).doNothing().when(flowFileRepository).updateRepository(any()); + final StatelessFlowTask boundaryTask = createTask(deferredRepository); + + assertThrows(IOException.class, () -> boundaryTask.completeInvocations(List.of(invocation), statelessProvRepo)); + + final List outputRecords = boundaryTask.getOutputRepositoryRecords(); + final ContentClaim firstBackingClaim = outputRecords.get(0).getCurrentClaim(); + final ContentClaim secondBackingClaim = outputRecords.get(1).getCurrentClaim(); + assertEquals(1, resourceClaimManager.getClaimantCount(firstBackingClaim.getResourceClaim())); + assertEquals(1, resourceClaimManager.getClaimantCount(secondBackingClaim.getResourceClaim())); + + deferredRepository.decrementClaimantCount(firstClaim); + deferredRepository.decrementClaimantCount(secondClaim); + deferredRepository.purge(); + + assertEquals(0, resourceClaimManager.getClaimantCount(firstBackingClaim.getResourceClaim())); + assertEquals(0, resourceClaimManager.getClaimantCount(secondBackingClaim.getResourceClaim())); + } + + @Test + public void testCreateOutputRecordsOnSuccessWithOnePortTwoConnections() throws IOException { final FlowFileQueue queue1 = mock(FlowFileQueue.class); final Connection conn1 = mock(Connection.class); when(conn1.getFlowFileQueue()).thenReturn(queue1); @@ -286,7 +375,7 @@ public void testCreateOutputRecordsOnSuccessWithOnePortTwoConnections() { } @Test - public void testCreateOutputRecordsOnSuccessWithTwoPorts() { + public void testCreateOutputRecordsOnSuccessWithTwoPorts() throws IOException { final Connection successConnection = mockConnection(); when(successPort.getConnections()).thenReturn(Collections.singleton(successConnection)); @@ -352,7 +441,7 @@ public void testCreateOutputRecordsOnSuccessWithTwoPorts() { } @Test - public void testUpdateClaimantCountSingleOutput() { + public void testUpdateClaimantCountSingleOutput() throws IOException { final Connection conn1 = mockConnection(); final Set singleConnectionSet = Collections.singleton(conn1); @@ -368,7 +457,7 @@ public void testUpdateClaimantCountSingleOutput() { } @Test - public void testUpdateClaimantCountWithClone() { + public void testUpdateClaimantCountWithClone() throws IOException { final Connection conn1 = mockConnection(); final Connection conn2 = mockConnection(); @@ -389,7 +478,7 @@ public void testUpdateClaimantCountWithClone() { } @Test - public void testContentClaimCountWhenMultipleFlowFilesTransferredToPort() { + public void testContentClaimCountWhenMultipleFlowFilesTransferredToPort() throws IOException { final Connection conn1 = mockConnection(); final Set singleConnectionSet = Collections.singleton(conn1); @@ -412,6 +501,26 @@ private Connection mockConnection() { return connection; } + private DeferredStatelessContentRepository createDeferredContentRepository() throws IOException { + final ByteArrayContentRepository backingRepository = new ByteArrayContentRepository(); + backingRepository.initialize(new StandardContentRepositoryContext(resourceClaimManager, EventReporter.NO_OP)); + final ProcessGroup processGroup = mock(ProcessGroup.class); + when(processGroup.resolveStatelessContentMaxHeap()).thenReturn(1024L); + return new DeferredStatelessContentRepository(processGroup, backingRepository, flowFileRepository, resourceClaimManager, EventReporter.NO_OP); + } + + private StatelessFlowTask createTask(final ContentRepository contentRepository) { + return new StatelessFlowTask.Builder() + .statelessFlow(statelessFlow) + .statelessGroupNode(statelessGroupNode) + .nifiFlowFileRepository(flowFileRepository) + .nifiContentRepository(contentRepository) + .nifiProvenanceRepository(provenanceRepository) + .flowFileEventRepository(flowFileEventRepository) + .logger(logger) + .build(); + } + @Test public void testProvenanceEventsCopiedFromStatelessFlow() { for (int i = 0; i < 5; i++) { diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/ProcessGroupResource.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/ProcessGroupResource.java index 6074142e03a5..25c3eac431d9 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/ProcessGroupResource.java +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/ProcessGroupResource.java @@ -71,6 +71,7 @@ import org.apache.nifi.flow.VersionedPropertyDescriptor; import org.apache.nifi.groups.VersionedComponentAdditions; import org.apache.nifi.parameter.ParameterContext; +import org.apache.nifi.processor.DataUnit; import org.apache.nifi.registry.flow.FlowRegistryBucket; import org.apache.nifi.registry.flow.FlowSnapshotContainer; import org.apache.nifi.registry.flow.RegisteredFlow; @@ -580,6 +581,27 @@ public Response updateProcessGroup( throw new IllegalArgumentException("Illegal value proposed for Max Concurrent Tasks: " + maxConcurrentTasks); } + final String statelessFlowFileContentInMemoryMax = requestProcessGroupDTO.getStatelessFlowFileContentInMemoryMax(); + if (statelessFlowFileContentInMemoryMax != null && !statelessFlowFileContentInMemoryMax.isBlank()) { + try { + DataUnit.parseDataSize(statelessFlowFileContentInMemoryMax.trim(), DataUnit.B); + } catch (final Exception e) { + throw new IllegalArgumentException("Illegal value proposed for Max In-Memory FlowFile Content: " + statelessFlowFileContentInMemoryMax); + } + } + + final String heapPercentage = requestProcessGroupDTO.getStatelessFlowFileContentInMemoryHeapPercentage(); + if (heapPercentage != null && !heapPercentage.isBlank()) { + try { + final int parsedHeapPercentage = Integer.parseInt(heapPercentage.trim()); + if (parsedHeapPercentage <= 0 || parsedHeapPercentage > 90) { + throw new IllegalArgumentException("Max In-Memory Heap Percentage must be greater than 0 and no more than 90: " + heapPercentage); + } + } catch (final NumberFormatException e) { + throw new IllegalArgumentException("Max In-Memory Heap Percentage must be a whole number: " + heapPercentage); + } + } + if (isReplicateRequest()) { return replicate(HttpMethod.PUT, requestProcessGroupEntity); } else if (isDisconnectedFromCluster()) { diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/dto/DtoFactory.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/dto/DtoFactory.java index 967abbd5a7a0..cf804fa631cf 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/dto/DtoFactory.java +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/dto/DtoFactory.java @@ -2840,6 +2840,9 @@ private ProcessGroupDTO createConciseProcessGroupDto(final ProcessGroup group) { dto.setResolvedExecutionEngine(group.resolveExecutionEngine().name()); dto.setMaxConcurrentTasks(group.getMaxConcurrentTasks()); dto.setStatelessFlowTimeout(group.getStatelessFlowTimeout()); + dto.setStatelessFlowFileContentInMemoryMax(group.getStatelessContentMaxHeap()); + final Integer heapPercentage = group.getStatelessContentMaxHeapPercentage(); + dto.setStatelessFlowFileContentInMemoryHeapPercentage(heapPercentage == null ? "" : Integer.toString(heapPercentage)); final ParameterContext parameterContext = group.getParameterContext(); if (parameterContext != null) { @@ -4832,6 +4835,8 @@ public ProcessGroupDTO copy(final ProcessGroupDTO original, final boolean deep) copy.setResolvedExecutionEngine(original.getResolvedExecutionEngine()); copy.setMaxConcurrentTasks(original.getMaxConcurrentTasks()); copy.setStatelessFlowTimeout(original.getStatelessFlowTimeout()); + copy.setStatelessFlowFileContentInMemoryMax(original.getStatelessFlowFileContentInMemoryMax()); + copy.setStatelessFlowFileContentInMemoryHeapPercentage(original.getStatelessFlowFileContentInMemoryHeapPercentage()); copy.setRunningCount(original.getRunningCount()); copy.setStoppedCount(original.getStoppedCount()); diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/dao/impl/StandardProcessGroupDAO.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/dao/impl/StandardProcessGroupDAO.java index 214591cd6941..e31ca1cfcae0 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/dao/impl/StandardProcessGroupDAO.java +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/dao/impl/StandardProcessGroupDAO.java @@ -110,6 +110,12 @@ public ProcessGroup createProcessGroup(String parentGroupId, ProcessGroupDTO pro if (processGroup.getExecutionEngine() != null) { group.setExecutionEngine(ExecutionEngine.valueOf(processGroup.getExecutionEngine())); } + if (processGroup.getStatelessFlowFileContentInMemoryMax() != null) { + group.setStatelessContentMaxHeap(processGroup.getStatelessFlowFileContentInMemoryMax()); + } + if (processGroup.getStatelessFlowFileContentInMemoryHeapPercentage() != null) { + group.setStatelessContentMaxHeapPercentage(processGroup.toStatelessFlowFileContentInMemoryHeapPercentage()); + } // add the process group group.setParent(parentGroup); @@ -133,6 +139,13 @@ public void verifyUpdate(final ProcessGroupDTO processGroup) { group.verifyCanSetExecutionEngine(ExecutionEngine.valueOf(executionEngine)); } + if (processGroup.getStatelessFlowFileContentInMemoryMax() != null) { + group.verifyCanSetStatelessContentMaxHeap(processGroup.getStatelessFlowFileContentInMemoryMax()); + } + if (processGroup.getStatelessFlowFileContentInMemoryHeapPercentage() != null) { + group.verifyCanSetStatelessContentMaxHeapPercentage(processGroup.toStatelessFlowFileContentInMemoryHeapPercentage()); + } + final VersionControlInformationDTO versionControlInfoDTO = processGroup.getVersionControlInformation(); final VersionControlInformation versionControlInformation = group.getVersionControlInformation(); if (versionControlInfoDTO != null) { @@ -497,6 +510,12 @@ public ProcessGroup updateProcessGroup(ProcessGroupDTO processGroupDTO) { if (processGroupDTO.getStatelessFlowTimeout() != null) { group.setStatelessFlowTimeout(processGroupDTO.getStatelessFlowTimeout()); } + if (processGroupDTO.getStatelessFlowFileContentInMemoryMax() != null) { + group.setStatelessContentMaxHeap(processGroupDTO.getStatelessFlowFileContentInMemoryMax()); + } + if (processGroupDTO.getStatelessFlowFileContentInMemoryHeapPercentage() != null) { + group.setStatelessContentMaxHeapPercentage(processGroupDTO.toStatelessFlowFileContentInMemoryHeapPercentage()); + } if (logFileSuffix != null) { group.setLogFileSuffix(logFileSuffix); diff --git a/nifi-frontend/src/main/frontend/apps/nifi/src/app/ui/common/component-dialogs/edit-process-group/edit-process-group.component.html b/nifi-frontend/src/main/frontend/apps/nifi/src/app/ui/common/component-dialogs/edit-process-group/edit-process-group.component.html index a11039a16682..b40684426078 100644 --- a/nifi-frontend/src/main/frontend/apps/nifi/src/app/ui/common/component-dialogs/edit-process-group/edit-process-group.component.html +++ b/nifi-frontend/src/main/frontend/apps/nifi/src/app/ui/common/component-dialogs/edit-process-group/edit-process-group.component.html @@ -137,6 +137,42 @@

{{ readonly ? 'Process Group Details' : 'Edit Process Group [readonly]="readonly" /> +
+ + + Max In-Memory FlowFile Content + + + + +
+
+ + + Max In-Memory Heap Percentage + + + + +
}
diff --git a/nifi-frontend/src/main/frontend/apps/nifi/src/app/ui/common/component-dialogs/edit-process-group/edit-process-group.component.spec.ts b/nifi-frontend/src/main/frontend/apps/nifi/src/app/ui/common/component-dialogs/edit-process-group/edit-process-group.component.spec.ts index 32a71776fa48..c6365cc30721 100644 --- a/nifi-frontend/src/main/frontend/apps/nifi/src/app/ui/common/component-dialogs/edit-process-group/edit-process-group.component.spec.ts +++ b/nifi-frontend/src/main/frontend/apps/nifi/src/app/ui/common/component-dialogs/edit-process-group/edit-process-group.component.spec.ts @@ -117,7 +117,9 @@ describe('EditProcessGroup', () => { }, executionEngine: 'INHERITED', maxConcurrentTasks: 1, - statelessFlowTimeout: '1 min' + statelessFlowTimeout: '1 min', + statelessFlowFileContentInMemoryMax: '0 B', + statelessGroupScheduledState: 'STOPPED' } } }; @@ -149,6 +151,85 @@ describe('EditProcessGroup', () => { expect(component).toBeTruthy(); }); + it('validates the stateless in-memory content maximum as a data size', () => { + component.executionEngineChanged('STATELESS'); + const control = component.editProcessGroupForm.get('statelessFlowFileContentInMemoryMax'); + + control?.setValue('not a size'); + expect(control?.valid).toBeFalsy(); + + control?.setValue('0.9 B'); + expect(control?.valid).toBeFalsy(); + + control?.setValue('1.5 KB'); + expect(control?.valid).toBeTruthy(); + + control?.setValue('999999999999999999999999999999999999999999999 TB'); + expect(control?.valid).toBeFalsy(); + + control?.setValue('100 MB'); + expect(control?.valid).toBeTruthy(); + + control?.setValue(' 100 MB '); + expect(control?.valid).toBeTruthy(); + + control?.setValue(''); + expect(control?.valid).toBeTruthy(); + + control?.setValue('0 B'); + expect(control?.valid).toBeTruthy(); + + control?.setValue('50%'); + expect(control?.valid).toBeFalsy(); + }); + + it('validates the stateless in-memory heap percentage', () => { + component.executionEngineChanged('STATELESS'); + const control = component.editProcessGroupForm.get('statelessFlowFileContentInMemoryHeapPercentage'); + + control?.setValue(0); + expect(control?.valid).toBeFalsy(); + + control?.setValue(1); + expect(control?.valid).toBeTruthy(); + + control?.setValue(90); + expect(control?.valid).toBeTruthy(); + + control?.setValue(''); + expect(control?.valid).toBeTruthy(); + + control?.setValue(91); + expect(control?.valid).toBeFalsy(); + + control?.setValue(-1); + expect(control?.valid).toBeFalsy(); + + control?.setValue(50.5); + expect(control?.valid).toBeFalsy(); + }); + + it('disables the stateless in-memory content maximum while the group is running', () => { + data.entity.component.statelessGroupScheduledState = 'RUNNING'; + const runningFixture = TestBed.createComponent(EditProcessGroup); + try { + const runningComponent = runningFixture.componentInstance; + runningComponent.executionEngineChanged('STATELESS'); + + expect( + runningComponent.editProcessGroupForm.get('statelessFlowFileContentInMemoryMax')?.disabled + ).toBeTruthy(); + expect( + runningComponent.editProcessGroupForm.get('statelessFlowFileContentInMemoryHeapPercentage') + ?.disabled + ).toBeTruthy(); + expect(runningComponent.editProcessGroupForm.get('statelessFlowTimeout')?.enabled).toBeTruthy(); + } finally { + data.entity.component.statelessGroupScheduledState = 'STOPPED'; + runningFixture.destroy(); + } + }); + it('verify parameter context value initialized', () => { expect(component.editProcessGroupForm.get('parameterContext')?.value).toEqual(selectedParameterContextId); }); @@ -218,7 +299,8 @@ describe('EditProcessGroup', () => { }, executionEngine: 'INHERITED', maxConcurrentTasks: 1, - statelessFlowTimeout: '1 min' + statelessFlowTimeout: '1 min', + statelessFlowFileContentInMemoryMax: '0 B' } } }; @@ -303,7 +385,8 @@ describe('EditProcessGroup', () => { }, executionEngine: 'INHERITED', maxConcurrentTasks: 1, - statelessFlowTimeout: '1 min' + statelessFlowTimeout: '1 min', + statelessFlowFileContentInMemoryMax: '0 B' } } }; diff --git a/nifi-frontend/src/main/frontend/apps/nifi/src/app/ui/common/component-dialogs/edit-process-group/edit-process-group.component.ts b/nifi-frontend/src/main/frontend/apps/nifi/src/app/ui/common/component-dialogs/edit-process-group/edit-process-group.component.ts index 05bf13a2dd16..0b6da9822bf3 100644 --- a/nifi-frontend/src/main/frontend/apps/nifi/src/app/ui/common/component-dialogs/edit-process-group/edit-process-group.component.ts +++ b/nifi-frontend/src/main/frontend/apps/nifi/src/app/ui/common/component-dialogs/edit-process-group/edit-process-group.component.ts @@ -17,7 +17,16 @@ import { Component, EventEmitter, Input, Output, inject } from '@angular/core'; import { MAT_DIALOG_DATA, MatDialogModule } from '@angular/material/dialog'; -import { FormBuilder, FormControl, FormGroup, FormsModule, ReactiveFormsModule, Validators } from '@angular/forms'; +import { + AbstractControl, + FormBuilder, + FormControl, + FormGroup, + FormsModule, + ReactiveFormsModule, + ValidationErrors, + Validators +} from '@angular/forms'; import { MatInputModule } from '@angular/material/input'; import { MatCheckboxModule } from '@angular/material/checkbox'; import { MatButtonModule } from '@angular/material/button'; @@ -69,6 +78,17 @@ import { ] }) export class EditProcessGroup extends TabbedDialog { + private static readonly DATA_SIZE_PATTERN = /^(\d+(?:\.\d+)?)\s*(B|KB|MB|GB|TB)$/i; + private static readonly MAXIMUM_HEAP_PERCENTAGE = 90; + private static readonly DATA_SIZE_MULTIPLIERS = { + B: 1n, + KB: 1n << 10n, + MB: 1n << 20n, + GB: 1n << 30n, + TB: 1n << 40n + }; + private static readonly MAXIMUM_DATA_SIZE_BYTES = 9223372036854775807n; + request = inject(MAT_DIALOG_DATA); private formBuilder = inject(FormBuilder); private client = inject(Client); @@ -125,6 +145,8 @@ export class EditProcessGroup extends TabbedDialog { protected readonly STATELESS: string = 'STATELESS'; private initialMaxConcurrentTasks: number; private initialStatelessFlowTimeout: string; + private initialStatelessFlowFileContentInMemoryMax: string; + private initialStatelessFlowFileContentInMemoryHeapPercentage: string | number; private _parameterContexts: ParameterContextEntity[] = []; editProcessGroupForm: FormGroup; @@ -226,6 +248,10 @@ export class EditProcessGroup extends TabbedDialog { this.initialMaxConcurrentTasks = request.entity.component.maxConcurrentTasks; this.initialStatelessFlowTimeout = request.entity.component.statelessFlowTimeout; + this.initialStatelessFlowFileContentInMemoryMax = + request.entity.component.statelessFlowFileContentInMemoryMax ?? ''; + this.initialStatelessFlowFileContentInMemoryHeapPercentage = + request.entity.component.statelessFlowFileContentInMemoryHeapPercentage ?? ''; this.executionEngineChanged(request.entity.component.executionEngine); } @@ -240,10 +266,78 @@ export class EditProcessGroup extends TabbedDialog { 'statelessFlowTimeout', new FormControl(this.initialStatelessFlowTimeout, Validators.required) ); + this.editProcessGroupForm.addControl( + 'statelessFlowFileContentInMemoryMax', + new FormControl( + { + value: this.initialStatelessFlowFileContentInMemoryMax, + disabled: this.request.entity.component.statelessGroupScheduledState !== 'STOPPED' + }, + EditProcessGroup.validateInMemoryContentMax + ) + ); + this.editProcessGroupForm.addControl( + 'statelessFlowFileContentInMemoryHeapPercentage', + new FormControl( + { + value: this.initialStatelessFlowFileContentInMemoryHeapPercentage, + disabled: this.request.entity.component.statelessGroupScheduledState !== 'STOPPED' + }, + EditProcessGroup.validateInMemoryHeapPercentage + ) + ); } else { this.editProcessGroupForm.removeControl('maxConcurrentTasks'); this.editProcessGroupForm.removeControl('statelessFlowTimeout'); + this.editProcessGroupForm.removeControl('statelessFlowFileContentInMemoryMax'); + this.editProcessGroupForm.removeControl('statelessFlowFileContentInMemoryHeapPercentage'); + } + } + + private static validateInMemoryContentMax(control: AbstractControl): ValidationErrors | null { + if (control.value === null || control.value === undefined || control.value === '') { + return null; + } + + if (typeof control.value !== 'string') { + return { dataSize: true }; + } + + const trimmedValue = control.value.trim(); + if (trimmedValue.length === 0) { + return null; } + + const match = EditProcessGroup.DATA_SIZE_PATTERN.exec(trimmedValue); + if (match === null) { + return { dataSize: true }; + } + + const size = match[1]; + const unit = match[2].toUpperCase() as keyof typeof EditProcessGroup.DATA_SIZE_MULTIPLIERS; + const [integerPart, fractionalPart = ''] = size.split('.'); + const unscaledSize = BigInt(integerPart + fractionalPart); + const divisor = 10n ** BigInt(fractionalPart.length); + const unscaledBytes = unscaledSize * EditProcessGroup.DATA_SIZE_MULTIPLIERS[unit]; + + if (unscaledBytes % divisor !== 0n || unscaledBytes / divisor > EditProcessGroup.MAXIMUM_DATA_SIZE_BYTES) { + return { dataSize: true }; + } + + return null; + } + + private static validateInMemoryHeapPercentage(control: AbstractControl): ValidationErrors | null { + if (control.value === null || control.value === undefined || control.value === '') { + return null; + } + + const percentage = Number(control.value); + if (!Number.isInteger(percentage) || percentage <= 0 || percentage > EditProcessGroup.MAXIMUM_HEAP_PERCENTAGE) { + return { heapPercentage: true }; + } + + return null; } submitForm() { @@ -279,6 +373,15 @@ export class EditProcessGroup extends TabbedDialog { if (this.editProcessGroupForm.get('executionEngine')?.value === this.STATELESS) { payload.component.maxConcurrentTasks = this.editProcessGroupForm.get('maxConcurrentTasks')?.value; payload.component.statelessFlowTimeout = this.editProcessGroupForm.get('statelessFlowTimeout')?.value; + payload.component.statelessFlowFileContentInMemoryMax = + this.editProcessGroupForm.get('statelessFlowFileContentInMemoryMax')?.value ?? ''; + const heapPercentage = this.editProcessGroupForm.get( + 'statelessFlowFileContentInMemoryHeapPercentage' + )?.value; + payload.component.statelessFlowFileContentInMemoryHeapPercentage = + heapPercentage === null || heapPercentage === undefined || heapPercentage === '' + ? '' + : String(heapPercentage); } this.editProcessGroup.next(payload); diff --git a/nifi-registry/nifi-registry-core/nifi-registry-flow-diff/src/main/java/org/apache/nifi/registry/flow/diff/DifferenceType.java b/nifi-registry/nifi-registry-core/nifi-registry-flow-diff/src/main/java/org/apache/nifi/registry/flow/diff/DifferenceType.java index cd62759da5a2..fe02ec36ed00 100644 --- a/nifi-registry/nifi-registry-core/nifi-registry-flow-diff/src/main/java/org/apache/nifi/registry/flow/diff/DifferenceType.java +++ b/nifi-registry/nifi-registry-core/nifi-registry-flow-diff/src/main/java/org/apache/nifi/registry/flow/diff/DifferenceType.java @@ -103,6 +103,16 @@ public enum DifferenceType { */ TIMEOUT_CHANGED("Timeout Changed"), + /** + * The Process Group has a different maximum amount of in-memory FlowFile content for the Stateless Engine + */ + STATELESS_IN_MEMORY_CONTENT_MAX_CHANGED("Stateless In-Memory Content Maximum Changed"), + + /** + * The Process Group has a different maximum in-memory FlowFile content heap percentage for the Stateless Engine + */ + STATELESS_IN_MEMORY_HEAP_PERCENTAGE_CHANGED("Stateless In-Memory Heap Percentage Changed"), + /** * The component has a different run schedule in each of the flows */ diff --git a/nifi-registry/nifi-registry-core/nifi-registry-flow-diff/src/main/java/org/apache/nifi/registry/flow/diff/StandardFlowComparator.java b/nifi-registry/nifi-registry-core/nifi-registry-flow-diff/src/main/java/org/apache/nifi/registry/flow/diff/StandardFlowComparator.java index 021185369532..7b6eabce285c 100644 --- a/nifi-registry/nifi-registry-core/nifi-registry-flow-diff/src/main/java/org/apache/nifi/registry/flow/diff/StandardFlowComparator.java +++ b/nifi-registry/nifi-registry-core/nifi-registry-flow-diff/src/main/java/org/apache/nifi/registry/flow/diff/StandardFlowComparator.java @@ -613,6 +613,8 @@ private void extractPGConfigDifferences(final VersionedProcessGroup groupA, fina addIfDifferent(differences, DifferenceType.SCHEDULED_STATE_CHANGED, groupA, groupB, VersionedProcessGroup::getScheduledState, true, ScheduledState.ENABLED); addIfDifferent(differences, DifferenceType.CONCURRENT_TASKS_CHANGED, groupA, groupB, VersionedProcessGroup::getMaxConcurrentTasks, true, 1); addIfDifferent(differences, DifferenceType.TIMEOUT_CHANGED, groupA, groupB, VersionedProcessGroup::getStatelessFlowTimeout, false, "1 min"); + addIfDifferent(differences, DifferenceType.STATELESS_IN_MEMORY_CONTENT_MAX_CHANGED, groupA, groupB, VersionedProcessGroup::getStatelessFlowFileContentInMemoryMax, false, null); + addIfDifferent(differences, DifferenceType.STATELESS_IN_MEMORY_HEAP_PERCENTAGE_CHANGED, groupA, groupB, VersionedProcessGroup::getStatelessFlowFileContentInMemoryHeapPercentage, true, null); } private void extractPGComponentsDifferences(final VersionedProcessGroup groupA, final VersionedProcessGroup groupB, final Set differences) { diff --git a/nifi-stateless/nifi-stateless-bundle/nifi-stateless-engine/src/main/java/org/apache/nifi/stateless/engine/StandardExecutionProgress.java b/nifi-stateless/nifi-stateless-bundle/nifi-stateless-engine/src/main/java/org/apache/nifi/stateless/engine/StandardExecutionProgress.java index 10a237831d46..435f905ecd9f 100644 --- a/nifi-stateless/nifi-stateless-bundle/nifi-stateless-engine/src/main/java/org/apache/nifi/stateless/engine/StandardExecutionProgress.java +++ b/nifi-stateless/nifi-stateless-bundle/nifi-stateless-engine/src/main/java/org/apache/nifi/stateless/engine/StandardExecutionProgress.java @@ -53,6 +53,7 @@ import java.util.Set; import java.util.concurrent.BlockingQueue; import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Consumer; public class StandardExecutionProgress implements ExecutionProgress { @@ -210,6 +211,7 @@ private TriggerResult createResult(final Runnable onAcknowledge, final Consumer< for (final String failurePortName : failurePortNames) { final List flowFilesForPort = outputFlowFiles.get(failurePortName); if (flowFilesForPort != null && !flowFilesForPort.isEmpty()) { + decrementOutputClaimantCounts(outputFlowFiles); throw new FailurePortEncounteredException("FlowFile was transferred to Port " + failurePortName + ", which is marked as a Failure Port", failurePortName); } } @@ -218,6 +220,7 @@ private TriggerResult createResult(final Runnable onAcknowledge, final Consumer< return new TriggerResult() { private volatile Throwable abortCause = null; + private final AtomicBoolean outputClaimantsReleased = new AtomicBoolean(false); @Override public boolean isSuccessful() { @@ -286,8 +289,8 @@ public byte[] readContentAsByteArray(final FlowFile flowFile) throws IOException public void acknowledge() { commitTracker.triggerCallbacks(); stateManagerProvider.commitUpdates(); + releaseOutputClaimants(); completionActionQueue.offer(CompletionAction.COMPLETE); - contentRepository.purge(); if (onAcknowledge != null) { onAcknowledge.run(); @@ -297,12 +300,20 @@ public void acknowledge() { @Override public void abort(final Throwable cause) { abortCause = new DataflowAbortedException("Dataflow was aborted", cause); + releaseOutputClaimants(); notifyExecutionFailed(abortCause); if (onFailure != null) { onFailure.accept(cause); } } + + private void releaseOutputClaimants() { + if (outputClaimantsReleased.compareAndSet(false, true)) { + decrementOutputClaimantCounts(outputFlowFiles); + contentRepository.purge(); + } + } }; } @@ -367,12 +378,17 @@ private List drainOutputQueues(final Port port) { final List flowFileRecords = new ArrayList<>(drainableQueue.size().getObjectCount()); drainableQueue.drainTo(flowFileRecords); portFlowFiles.addAll(flowFileRecords); + } - for (final FlowFileRecord flowFileRecord : flowFileRecords) { + return portFlowFiles; + } + + private void decrementOutputClaimantCounts(final Map> outputFlowFiles) { + for (final List flowFiles : outputFlowFiles.values()) { + for (final FlowFile flowFile : flowFiles) { + final FlowFileRecord flowFileRecord = (FlowFileRecord) flowFile; contentRepository.decrementClaimantCount(flowFileRecord.getContentClaim()); } } - - return portFlowFiles; } } diff --git a/nifi-stateless/nifi-stateless-bundle/nifi-stateless-engine/src/test/java/org/apache/nifi/stateless/engine/StandardExecutionProgressTest.java b/nifi-stateless/nifi-stateless-bundle/nifi-stateless-engine/src/test/java/org/apache/nifi/stateless/engine/StandardExecutionProgressTest.java new file mode 100644 index 000000000000..9ecb84e2ef4c --- /dev/null +++ b/nifi-stateless/nifi-stateless-bundle/nifi-stateless-engine/src/test/java/org/apache/nifi/stateless/engine/StandardExecutionProgressTest.java @@ -0,0 +1,131 @@ +/* + * 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.nifi.stateless.engine; + +import org.apache.nifi.components.state.StatelessStateManagerProvider; +import org.apache.nifi.connectable.Connection; +import org.apache.nifi.connectable.Port; +import org.apache.nifi.controller.queue.QueueSize; +import org.apache.nifi.controller.repository.ContentRepository; +import org.apache.nifi.controller.repository.FlowFileRecord; +import org.apache.nifi.controller.repository.claim.ContentClaim; +import org.apache.nifi.groups.ProcessGroup; +import org.apache.nifi.stateless.flow.DataflowTriggerContext; +import org.apache.nifi.stateless.flow.FailurePortEncounteredException; +import org.apache.nifi.stateless.flow.TriggerResult; +import org.apache.nifi.stateless.queue.DrainableFlowFileQueue; +import org.apache.nifi.stateless.repository.RepositoryContextFactory; +import org.apache.nifi.stateless.session.AsynchronousCommitTracker; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.List; +import java.util.Set; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingQueue; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class StandardExecutionProgressTest { + + @Test + void testOutputClaimHeldUntilAcknowledged() { + final TestContext context = createTestContext(); + + context.progress().enqueueTriggerResult(() -> { }, failure -> { }); + final TriggerResult result = context.results().remove(); + + verify(context.contentRepository(), never()).decrementClaimantCount(context.contentClaim()); + + result.acknowledge(); + + verify(context.contentRepository()).decrementClaimantCount(context.contentClaim()); + verify(context.contentRepository()).purge(); + } + + @Test + void testRepeatedAbortReleasesOutputClaimOnce() { + final TestContext context = createTestContext(); + context.progress().enqueueTriggerResult(() -> { }, failure -> { }); + final TriggerResult result = context.results().remove(); + + result.abort(new IOException("Processing failed")); + result.abort(new IOException("Processing failed again")); + + verify(context.contentRepository()).decrementClaimantCount(context.contentClaim()); + verify(context.contentRepository()).purge(); + verify(context.purgeAction()).purge(); + } + + @Test + void testFailurePortReleasesOutputClaimOnce() { + final TestContext context = createTestContext(Set.of("out")); + + assertThrows(FailurePortEncounteredException.class, () -> context.progress().enqueueTriggerResult(() -> { }, failure -> { })); + context.progress().notifyExecutionFailed(new IOException("Failure port encountered")); + + verify(context.contentRepository(), times(1)).decrementClaimantCount(context.contentClaim()); + verify(context.purgeAction()).purge(); + } + + private TestContext createTestContext() { + return createTestContext(Set.of()); + } + + private TestContext createTestContext(final Set failurePortNames) { + final ContentClaim contentClaim = mock(ContentClaim.class); + final FlowFileRecord flowFile = mock(FlowFileRecord.class); + when(flowFile.getContentClaim()).thenReturn(contentClaim); + + final DrainableFlowFileQueue flowFileQueue = mock(DrainableFlowFileQueue.class); + when(flowFileQueue.size()).thenReturn(new QueueSize(1, 1L)); + doAnswer(invocation -> { + final List destination = invocation.getArgument(0); + destination.add(flowFile); + return null; + }).when(flowFileQueue).drainTo(anyList()); + + final Connection connection = mock(Connection.class); + when(connection.getFlowFileQueue()).thenReturn(flowFileQueue); + final Port outputPort = mock(Port.class); + when(outputPort.getName()).thenReturn("out"); + when(outputPort.getIncomingConnections()).thenReturn(List.of(connection)); + final ProcessGroup rootGroup = mock(ProcessGroup.class); + when(rootGroup.getOutputPorts()).thenReturn(Set.of(outputPort)); + + final ContentRepository contentRepository = mock(ContentRepository.class); + final RepositoryContextFactory repositoryContextFactory = mock(RepositoryContextFactory.class); + when(repositoryContextFactory.getContentRepository()).thenReturn(contentRepository); + final BlockingQueue results = new LinkedBlockingQueue<>(); + final FlowPurgeAction purgeAction = mock(FlowPurgeAction.class); + final StandardExecutionProgress progress = new StandardExecutionProgress(rootGroup, List.of(), results, repositoryContextFactory, failurePortNames, + mock(AsynchronousCommitTracker.class), mock(StatelessStateManagerProvider.class), mock(DataflowTriggerContext.class), purgeAction); + return new TestContext(progress, contentRepository, purgeAction, contentClaim, results); + } + + private record TestContext(StandardExecutionProgress progress, ContentRepository contentRepository, FlowPurgeAction purgeAction, ContentClaim contentClaim, + BlockingQueue results) { + } +} diff --git a/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/java/org/apache/nifi/processors/tests/system/PassThrough.java b/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/java/org/apache/nifi/processors/tests/system/PassThrough.java index cc4fb4a31ab5..08a7e83bcdc2 100644 --- a/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/java/org/apache/nifi/processors/tests/system/PassThrough.java +++ b/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/java/org/apache/nifi/processors/tests/system/PassThrough.java @@ -17,23 +17,39 @@ package org.apache.nifi.processors.tests.system; +import org.apache.nifi.components.PropertyDescriptor; import org.apache.nifi.flowfile.FlowFile; import org.apache.nifi.processor.AbstractProcessor; import org.apache.nifi.processor.ProcessContext; import org.apache.nifi.processor.ProcessSession; import org.apache.nifi.processor.Relationship; import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processor.util.StandardValidators; +import java.io.File; import java.util.Collections; +import java.util.List; import java.util.Set; public class PassThrough extends AbstractProcessor { + private static final PropertyDescriptor GATE_FILE = new PropertyDescriptor.Builder() + .name("Gate File") + .description("An optional file path. FlowFiles are transferred only when this file exists.") + .required(false) + .addValidator(StandardValidators.NON_BLANK_VALIDATOR) + .build(); + private static final Relationship REL_SUCCESS = new Relationship.Builder() .name("success") .description("Everything goes here") .build(); + @Override + protected List getSupportedPropertyDescriptors() { + return List.of(GATE_FILE); + } + @Override public Set getRelationships() { return Collections.singleton(REL_SUCCESS); @@ -41,7 +57,13 @@ public Set getRelationships() { @Override public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException { - FlowFile flowFile = session.get(); + final String gateFilePath = context.getProperty(GATE_FILE).getValue(); + if (gateFilePath != null && !new File(gateFilePath).exists()) { + context.yield(); + return; + } + + final FlowFile flowFile = session.get(); if (flowFile == null) { return; } diff --git a/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/NiFiClientUtil.java b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/NiFiClientUtil.java index faef5b9a5245..b9324346c8f4 100644 --- a/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/NiFiClientUtil.java +++ b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/NiFiClientUtil.java @@ -1287,6 +1287,12 @@ public ProcessorEntity updateProcessorRunDuration(final ProcessorEntity currentE return updateProcessorConfig(currentEntity, config); } + public ProcessorEntity updateProcessorYieldDuration(final ProcessorEntity currentEntity, final String yieldDuration) throws NiFiClientException, IOException { + final ProcessorConfigDTO config = new ProcessorConfigDTO(); + config.setYieldDuration(yieldDuration); + return updateProcessorConfig(currentEntity, config); + } + public ProcessorEntity updateProcessorSchedulingPeriod(final ProcessorEntity currentEntity, final String schedulingPeriod) throws NiFiClientException, IOException { final ProcessorConfigDTO config = new ProcessorConfigDTO(); config.setSchedulingPeriod(schedulingPeriod); @@ -2890,4 +2896,23 @@ public ProcessGroupEntity markStateless(final ProcessGroupEntity group, final St return nifiClient.getProcessGroupClient().updateProcessGroup(group); } + + public ProcessGroupEntity markStateless(final ProcessGroupEntity group, final String timeout, final String inMemoryContentMax) + throws NiFiClientException, IOException { + group.getComponent().setStatelessFlowTimeout(timeout); + group.getComponent().setExecutionEngine("STATELESS"); + group.getComponent().setStatelessFlowFileContentInMemoryMax(inMemoryContentMax); + group.getComponent().setStatelessFlowFileContentInMemoryHeapPercentage(""); + + return nifiClient.getProcessGroupClient().updateProcessGroup(group); + } + + public ProcessGroupEntity setStatelessFlowFileContentInMemoryMax(final ProcessGroupEntity group, final String inMemoryContentMax) + throws NiFiClientException, IOException { + final ProcessGroupEntity current = nifiClient.getProcessGroupClient().getProcessGroup(group.getId()); + current.getComponent().setStatelessFlowFileContentInMemoryMax(inMemoryContentMax); + current.getComponent().setStatelessFlowFileContentInMemoryHeapPercentage(""); + return nifiClient.getProcessGroupClient().updateProcessGroup(current); + } + } diff --git a/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/stateless/StatelessInMemoryContentIT.java b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/stateless/StatelessInMemoryContentIT.java new file mode 100644 index 000000000000..4522c8df568f --- /dev/null +++ b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/stateless/StatelessInMemoryContentIT.java @@ -0,0 +1,285 @@ +/* + * 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.nifi.tests.system.stateless; + +import jakarta.ws.rs.WebApplicationException; +import org.apache.nifi.tests.system.NiFiSystemIT; +import org.apache.nifi.toolkit.client.NiFiClientException; +import org.apache.nifi.web.api.entity.ConnectionEntity; +import org.apache.nifi.web.api.entity.PortEntity; +import org.apache.nifi.web.api.entity.ProcessGroupEntity; +import org.apache.nifi.web.api.entity.ProcessorEntity; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import java.util.stream.Stream; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; + +@Timeout(value = 30, unit = TimeUnit.SECONDS) +public class StatelessInMemoryContentIT extends NiFiSystemIT { + + private static final String HELLO_WORLD = "Hello World"; + private static final String EXCLAMATIONS = "!!!"; + + // Result of reversing "Hello World" to "dlroW olleH", appending "!!!", then reversing "dlroW olleH!!!" back. + private static final String TRANSFORMED = EXCLAMATIONS + HELLO_WORLD; + + private static final String LARGE_BUDGET = "1 MB"; + private static final String TINY_BUDGET = "10 B"; + private static final String GATE_YIELD_DURATION = "10 millis"; + + @Test + public void testContentUnderBudgetStaysInMemory() throws NiFiClientException, IOException, InterruptedException { + final SelfContainedFlow flow = createSelfContainedTransformFlow(LARGE_BUDGET); + final Map contentFileSizesBeforeStart = getContentFileSizes(); + + getClientUtil().startProcessGroupComponents(flow.groupId()); + + waitFor(() -> Files.exists(flow.markerFile())); + assertFalse(hasContentRepositoryGrown(contentFileSizesBeforeStart)); + Files.createFile(flow.gateFile()); + waitFor(() -> getProcessorFlowFilesIn(flow.matchedTerminateId()) >= 1); + getClientUtil().stopProcessGroupComponents(flow.groupId()); + + assertEquals(0, getProcessorFlowFilesIn(flow.unmatchedTerminateId())); + } + + @Test + public void testContentOverBudgetSpills() throws NiFiClientException, IOException, InterruptedException { + final SelfContainedFlow flow = createSelfContainedTransformFlow(TINY_BUDGET); + final Map contentFileSizesBeforeStart = getContentFileSizes(); + + getClientUtil().startProcessGroupComponents(flow.groupId()); + + waitFor(() -> Files.exists(flow.markerFile())); + waitFor(() -> hasContentRepositoryGrown(contentFileSizesBeforeStart)); + Files.createFile(flow.gateFile()); + waitFor(() -> getProcessorFlowFilesIn(flow.matchedTerminateId()) >= 1); + getClientUtil().stopProcessGroupComponents(flow.groupId()); + + assertEquals(0, getProcessorFlowFilesIn(flow.unmatchedTerminateId())); + } + + @Test + public void testInputOutputPortsUnderBudget() throws NiFiClientException, IOException, InterruptedException { + verifyInputOutputPortsReverseContent(LARGE_BUDGET, HELLO_WORLD); + } + + @Test + public void testInputOutputPortsSpillOver() throws NiFiClientException, IOException, InterruptedException { + final StringBuilder builder = new StringBuilder(); + for (int i = 0; i < 200; i++) { + builder.append("ABCDEFGHIJ"); + } + + verifyInputOutputPortsReverseContent(TINY_BUDGET, builder.toString()); + } + + @Test + public void testCannotChangeInMemoryMaxWhileRunning() throws NiFiClientException, IOException, InterruptedException { + final ProcessGroupEntity statelessGroup = getClientUtil().createProcessGroup("Stateless", "root"); + getClientUtil().markStateless(statelessGroup, "1 min", LARGE_BUDGET); + final String groupId = statelessGroup.getId(); + + final ProcessorEntity generate = getClientUtil().createProcessor(GENERATE_FLOWFILE, groupId); + getClientUtil().updateProcessorProperties(generate, Map.of("Text", HELLO_WORLD)); + final ProcessorEntity terminate = getClientUtil().createProcessor(TERMINATE_FLOWFILE, groupId); + getClientUtil().createConnection(generate, terminate, SUCCESS, groupId); + + getClientUtil().waitForValidProcessor(generate.getId()); + getClientUtil().startProcessGroupComponents(groupId); + + waitFor(() -> getProcessorFlowFilesIn(terminate.getId()) >= 1); + + // The in-memory budget is applied when the Stateless flow starts, so it cannot be changed while the group is running. + final NiFiClientException exception = assertThrows( + NiFiClientException.class, () -> getClientUtil().setStatelessFlowFileContentInMemoryMax(statelessGroup, TINY_BUDGET)); + final WebApplicationException cause = assertInstanceOf(WebApplicationException.class, exception.getCause()); + assertEquals(409, cause.getResponse().getStatus()); + + getClientUtil().stopProcessGroupComponents(groupId); + + final ProcessGroupEntity stoppedGroup = getClientUtil().setStatelessFlowFileContentInMemoryMax(statelessGroup, TINY_BUDGET); + assertEquals(TINY_BUDGET, stoppedGroup.getComponent().getStatelessFlowFileContentInMemoryMax()); + + final Map contentFileSizesBeforeRestart = getContentFileSizes(); + getClientUtil().startProcessGroupComponents(groupId); + waitFor(() -> getProcessorFlowFilesIn(terminate.getId()) >= 2); + waitFor(() -> hasContentRepositoryGrown(contentFileSizesBeforeRestart)); + getClientUtil().stopProcessGroupComponents(groupId); + } + + /** + * Builds a flow in which a FlowFile is generated outside the Stateless group, sent into it through an Input Port, has its content rewritten inside the group by + * ReverseContents, and leaves through an Output Port into a connection whose destination is never started. The FlowFile therefore remains queued outside the + * group, which keeps the content that left the group referenced in the Content Repository. The transformed content is verified byte-for-byte. For the spill case, + * a file-controlled processor keeps the transformed FlowFile inside the Stateless group until its claim is confirmed in the Content Repository. + */ + private void verifyInputOutputPortsReverseContent(final String budget, final String content) throws NiFiClientException, IOException, InterruptedException { + final ProcessorEntity generate = getClientUtil().createProcessor(GENERATE_FLOWFILE); + getClientUtil().updateProcessorProperties(generate, Map.of("Text", content)); + final boolean verifySpill = TINY_BUDGET.equals(budget); + + final ProcessGroupEntity statelessGroup = getClientUtil().createProcessGroup("Stateless", "root"); + getClientUtil().markStateless(statelessGroup, "1 min", budget); + final String groupId = statelessGroup.getId(); + + final PortEntity inputPort = getClientUtil().createInputPort("In", groupId); + final PortEntity outputPort = getClientUtil().createOutputPort("Out", groupId); + + final ProcessorEntity reverse = getClientUtil().createProcessor(REVERSE_CONTENTS, groupId); + getClientUtil().createConnection(inputPort, reverse, groupId); + final Path gateFile; + if (verifySpill) { + gateFile = prepareGateFile("input-output-" + groupId); + final ProcessorEntity passThrough = getClientUtil().createProcessor("PassThrough", groupId); + getClientUtil().updateProcessorProperties(passThrough, Map.of("Gate File", gateFile.toString())); + getClientUtil().updateProcessorYieldDuration(passThrough, GATE_YIELD_DURATION); + getClientUtil().createConnection(reverse, passThrough, SUCCESS, groupId); + getClientUtil().createConnection(passThrough, outputPort, SUCCESS); + getClientUtil().waitForValidProcessor(passThrough.getId()); + } else { + gateFile = null; + getClientUtil().createConnection(reverse, outputPort, SUCCESS); + } + + final ProcessorEntity terminate = getClientUtil().createProcessor(TERMINATE_FLOWFILE); + final ConnectionEntity inputToStateless = getClientUtil().createConnection(generate, inputPort, SUCCESS); + final ConnectionEntity outputToTerminate = getClientUtil().createConnection(outputPort, terminate); + + getClientUtil().waitForValidProcessor(generate.getId()); + getClientUtil().waitForValidProcessor(reverse.getId()); + + getClientUtil().runProcessorOnce(generate); + waitForQueueCount(inputToStateless.getId(), 1); + final Map contentFileSizesBeforeStateless = getContentFileSizes(); + getClientUtil().startProcessGroupComponents(groupId); + + if (verifySpill) { + waitFor(() -> hasContentRepositoryGrown(contentFileSizesBeforeStateless)); + assertEquals(0, getConnectionQueueSize(outputToTerminate.getId())); + Files.createFile(gateFile); + } + + waitForQueueCount(outputToTerminate.getId(), 1); + + final String expected = new StringBuilder(content).reverse().toString(); + final String outputContent = getClientUtil().getFlowFileContentAsUtf8(outputToTerminate.getId(), 0); + assertEquals(expected, outputContent); + assertEquals(content.length(), getClientUtil().getQueueFlowFile(outputToTerminate.getId(), 0).getFlowFile().getSize()); + + getClientUtil().stopProcessGroupComponents(groupId); + } + + private SelfContainedFlow createSelfContainedTransformFlow(final String budget) throws NiFiClientException, IOException, InterruptedException { + final ProcessGroupEntity statelessGroup = getClientUtil().createProcessGroup("Stateless", "root"); + getClientUtil().markStateless(statelessGroup, "1 min", budget); + final String groupId = statelessGroup.getId(); + + final ProcessorEntity generate = getClientUtil().createProcessor(GENERATE_FLOWFILE, groupId); + getClientUtil().updateProcessorProperties(generate, Map.of("Text", HELLO_WORLD)); + + final ProcessorEntity reverseFirst = getClientUtil().createProcessor(REVERSE_CONTENTS, groupId); + final ProcessorEntity append = getClientUtil().createProcessor("UpdateContent", groupId); + getClientUtil().updateProcessorProperties(append, Map.of("Content", EXCLAMATIONS, "Update Strategy", "Append")); + final ProcessorEntity reverseSecond = getClientUtil().createProcessor(REVERSE_CONTENTS, groupId); + + final ProcessorEntity verify = getClientUtil().createProcessor("VerifyContents", groupId); + getClientUtil().updateProcessorProperties(verify, Map.of("matched", TRANSFORMED)); + final Path markerFile = new File(getNiFiInstance().getInstanceDirectory(), "target/stateless-content-marker-" + groupId).getAbsoluteFile().toPath(); + Files.deleteIfExists(markerFile); + final ProcessorEntity writeMarker = getClientUtil().createProcessor("WriteToFile", groupId); + getClientUtil().updateProcessorProperties(writeMarker, Map.of("Filename", markerFile.toString())); + getClientUtil().setAutoTerminatedRelationships(writeMarker, "failure"); + final Path gateFile = prepareGateFile("self-contained-" + groupId); + final ProcessorEntity passThrough = getClientUtil().createProcessor("PassThrough", groupId); + getClientUtil().updateProcessorProperties(passThrough, Map.of("Gate File", gateFile.toString())); + getClientUtil().updateProcessorYieldDuration(passThrough, GATE_YIELD_DURATION); + final ProcessorEntity matchedTerminate = getClientUtil().createProcessor(TERMINATE_FLOWFILE, groupId); + final ProcessorEntity unmatchedTerminate = getClientUtil().createProcessor(TERMINATE_FLOWFILE, groupId); + + getClientUtil().createConnection(generate, reverseFirst, SUCCESS, groupId); + getClientUtil().createConnection(reverseFirst, append, SUCCESS, groupId); + getClientUtil().createConnection(append, reverseSecond, SUCCESS, groupId); + getClientUtil().createConnection(reverseSecond, verify, SUCCESS, groupId); + getClientUtil().createConnection(verify, writeMarker, "matched", groupId); + getClientUtil().createConnection(writeMarker, passThrough, SUCCESS, groupId); + getClientUtil().createConnection(passThrough, matchedTerminate, SUCCESS, groupId); + getClientUtil().createConnection(verify, unmatchedTerminate, "unmatched", groupId); + + getClientUtil().waitForValidProcessor(generate.getId()); + getClientUtil().waitForValidProcessor(reverseFirst.getId()); + getClientUtil().waitForValidProcessor(append.getId()); + getClientUtil().waitForValidProcessor(reverseSecond.getId()); + getClientUtil().waitForValidProcessor(verify.getId()); + getClientUtil().waitForValidProcessor(writeMarker.getId()); + getClientUtil().waitForValidProcessor(passThrough.getId()); + + return new SelfContainedFlow(groupId, matchedTerminate.getId(), unmatchedTerminate.getId(), markerFile, gateFile); + } + + private Path prepareGateFile(final String name) throws IOException { + final Path gateFile = new File(getNiFiInstance().getInstanceDirectory(), "target/stateless-content-gate-" + name).getAbsoluteFile().toPath(); + Files.createDirectories(gateFile.getParent()); + Files.deleteIfExists(gateFile); + return gateFile; + } + + private int getProcessorFlowFilesIn(final String processorId) throws NiFiClientException, IOException { + return getNifiClient().getProcessorClient().getProcessor(processorId).getStatus().getAggregateSnapshot().getFlowFilesIn(); + } + + private Map getContentFileSizes() throws IOException { + final Map contentFileSizes = new HashMap<>(); + final File contentRepository = new File(getNiFiInstance().getInstanceDirectory(), "content_repository"); + if (!contentRepository.exists()) { + return contentFileSizes; + } + + try (final Stream paths = Files.walk(contentRepository.toPath())) { + paths.filter(Files::isRegularFile).forEach(path -> contentFileSizes.put(path, path.toFile().length())); + } + + return contentFileSizes; + } + + private boolean hasContentRepositoryGrown(final Map originalFileSizes) throws IOException { + final Map currentFileSizes = getContentFileSizes(); + for (final Map.Entry entry : currentFileSizes.entrySet()) { + if (entry.getValue() > originalFileSizes.getOrDefault(entry.getKey(), 0L)) { + return true; + } + } + + return false; + } + + private record SelfContainedFlow(String groupId, String matchedTerminateId, String unmatchedTerminateId, Path markerFile, Path gateFile) { + } +}