diff --git a/nifi-docs/src/main/asciidoc/user-guide.adoc b/nifi-docs/src/main/asciidoc/user-guide.adoc index b5c0eee2c966..2a95d8e6a765 100644 --- a/nifi-docs/src/main/asciidoc/user-guide.adoc +++ b/nifi-docs/src/main/asciidoc/user-guide.adoc @@ -1781,6 +1781,13 @@ 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 from `0` to `90`. A value of `0 B` or `0` percent means zero -- 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. The default is `0` percent and no data size, which uses no in-memory content. The configured maximum applies to the Stateless Process Group +as a whole, including all Concurrent Tasks. 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 on disk. + ==== Data Ordering @@ -1802,7 +1809,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-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 1a830d5420e2..641736a15af2 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 @@ -40,6 +40,8 @@ public class ProcessGroupDTO extends ComponentDTO { private String executionEngine; private Integer maxConcurrentTasks; private String statelessFlowTimeout; + private String statelessFlowFileContentInMemoryMax; + private String statelessFlowFileContentInMemoryHeapPercentage; private Integer runningCount; private Integer stoppedCount; @@ -422,4 +424,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, from 0 to 90. " + + "A value of 0 means zero percent of the heap. 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. The default is 0.") + 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("Illegal value proposed for Max In-Memory Heap Percentage: " + 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 600178a68661..1e426ef38f54 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,8 @@ 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 = 0; private volatile Authorizable explicitParentAuthorizable; private final FlowFileActivity flowFileActivity = new ProcessGroupFlowFileActivity(this); @@ -233,6 +239,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 +3768,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 +3794,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 +4789,128 @@ public void setStatelessFlowTimeout(final String statelessFlowTimeout) { } } + @Override + public String getStatelessContentMaxHeap() { + return statelessFlowFileContentInMemoryMax; + } + + @Override + public void setStatelessContentMaxHeap(final String maxSize) { + writeLock.lock(); + try { + verifyCanSetStatelessContentMaxHeap(maxSize); + this.statelessFlowFileContentInMemoryMax = normalizeStatelessContentMaxHeap(maxSize); + } finally { + writeLock.unlock(); + } + } + + @Override + public Integer getStatelessContentMaxHeapPercentage() { + return statelessFlowFileContentInMemoryHeapPercentage; + } + + @Override + public void setStatelessContentMaxHeapPercentage(final Integer heapPercentage) { + writeLock.lock(); + try { + verifyCanSetStatelessContentMaxHeapPercentage(heapPercentage); + this.statelessFlowFileContentInMemoryHeapPercentage = heapPercentage; + } finally { + writeLock.unlock(); + } + } + + @Override + public long resolveStatelessContentMaxHeap() { + return resolveStatelessContentMaxHeap(getStatelessContentMaxHeap(), getStatelessContentMaxHeapPercentage()); + } + + @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 between 0 and " + 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..007364ddaf7f 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; @@ -404,6 +408,127 @@ void shouldPropagateLoggingAttributesChangesToChildren() { assertEquals(expected, leaf.getLoggingAttributes()); } + @Test + void testStatelessContentMaxHeapDefaultsToZeroPercentAndUnsetSize() { + assertNull(processGroup.getStatelessContentMaxHeap()); + assertEquals(0, 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()); + } + + @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()); + + processGroup.setStatelessContentMaxHeapPercentage(0); + assertEquals(0, processGroup.getStatelessContentMaxHeapPercentage()); + assertEquals(0L, processGroup.resolveStatelessContentMaxHeap()); + } + + @Test + void testResolveStatelessContentMaxHeapUsesConfiguredLimitWhenTheOtherIsUnset() { + processGroup.setStatelessContentMaxHeapPercentage(null); + processGroup.setStatelessContentMaxHeap("2 MB"); + assertEquals(2L * 1024 * 1024, processGroup.resolveStatelessContentMaxHeap()); + + processGroup.setStatelessContentMaxHeap(null); + processGroup.setStatelessContentMaxHeapPercentage(25); + final long expectedTwentyFivePercent = new BigDecimal("25") + .multiply(BigDecimal.valueOf(Runtime.getRuntime().maxMemory())) + .divide(BigDecimal.valueOf(100), 0, RoundingMode.DOWN) + .longValue(); + assertEquals(expectedTwentyFivePercent, processGroup.resolveStatelessContentMaxHeap()); + } + + @Test + void testResolveStatelessContentMaxHeapUsesTheSmallerLimitWhenBothAreSet() { + processGroup.setStatelessContentMaxHeap("1 MB"); + processGroup.setStatelessContentMaxHeapPercentage(90); + assertEquals(1024L * 1024L, processGroup.resolveStatelessContentMaxHeap()); + + processGroup.setStatelessContentMaxHeap("100 GB"); + processGroup.setStatelessContentMaxHeapPercentage(50); + assertEquals(Runtime.getRuntime().maxMemory() / 2, processGroup.resolveStatelessContentMaxHeap()); + + processGroup.setStatelessContentMaxHeap("4 GB"); + processGroup.setStatelessContentMaxHeapPercentage(0); + assertEquals(0L, processGroup.resolveStatelessContentMaxHeap()); + } + + @Test + void testResolveStatelessContentMaxHeapWhenBothAreUnsetIsZero() { + processGroup.setStatelessContentMaxHeapPercentage(null); + processGroup.setStatelessContentMaxHeap(null); + assertNull(processGroup.getStatelessContentMaxHeap()); + assertNull(processGroup.getStatelessContentMaxHeapPercentage()); + assertEquals(0L, processGroup.resolveStatelessContentMaxHeap()); + } + + @Test + void testSetStatelessContentMaxHeapTreatsBlankAsUnset() { + processGroup.setStatelessContentMaxHeapPercentage(null); + processGroup.setStatelessContentMaxHeap("100 MB"); + + processGroup.setStatelessContentMaxHeap(" "); + assertNull(processGroup.getStatelessContentMaxHeap()); + assertEquals(0L, processGroup.resolveStatelessContentMaxHeap()); + } + + @Test + void testSetStatelessContentMaxHeapRejectsInvalidDataSize() { + processGroup.setStatelessContentMaxHeapPercentage(null); + 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(91)); + assertThrows(IllegalArgumentException.class, () -> processGroup.setStatelessContentMaxHeapPercentage(-1)); + } + + @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.setStatelessContentMaxHeap("1024 KB"); + assertEquals("1024 KB", runningGroup.getStatelessContentMaxHeap()); + assertThrows(IllegalStateException.class, () -> runningGroup.setStatelessContentMaxHeap("2 MB")); + assertThrows(IllegalStateException.class, () -> runningGroup.setStatelessContentMaxHeapPercentage(0)); + } + @Test void testFindOwningConnectorIdentifierWalksParentHierarchy() { final StandardProcessGroup connectorGroup = createStandardProcessGroup("connector-group", "connector-1"); 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..c8e991f5f234 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, from 0 to 90. A value of 0 means zero percent of the heap. A null value means this limit is not configured. The default is 0. + */ + 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, from 0 to 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, from 0 to 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..d48b8b33f954 --- /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 ContentRepository { + 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; + } + + /** + * Ensures that the content for the given claim is accessible outside of the Stateless Process Group and returns a Content Claim that references it in the + * NiFi instance's Content Repository. When FlowFile content was buffered in memory, this writes it to the NiFi Content Repository and returns the new claim. + * When content was not buffered in memory (or the Process Group is not configured to buffer content in memory), the claim is returned unchanged. + * + * @param claim the claim to make externally accessible + * @return a Content Claim whose content is stored in the NiFi instance's Content Repository + * @throws IOException if the content cannot be written to the NiFi Content Repository + */ + public ContentClaim exportForExternalUse(final ContentClaim claim) throws IOException { + final ContentRepository resolved = getDelegate(); + if (resolved instanceof final SpillableContentRepository spillableContentRepository) { + return spillableContentRepository.exportForExternalUse(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 #exportForExternalUse(ContentClaim)} + */ + public void commitExportForExternalUse(final ContentClaim claim) { + final ContentRepository resolved = getDelegate(); + if (resolved instanceof final SpillableContentRepository spillableContentRepository) { + spillableContentRepository.commitExportForExternalUse(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..30c66ca94ac5 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/repository/SpillableContentRepository.java @@ -0,0 +1,836 @@ +/* + * 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.commons.io.output.UnsynchronizedByteArrayOutputStream; +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.stream.io.StreamUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.file.Files; +import java.nio.file.OpenOption; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.HashSet; +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 (on-disk) 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 write checks the running total of buffered bytes and, if the write would exceed the configured + * size, the bytes buffered so far are flushed to the backing repository and the remainder of the claim is written there. + * + *

+ * 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. Exporting prepares the backing claim while + * retaining that ownership until {@link #commitExportForExternalUse(ContentClaim)} confirms that the NiFi FlowFile Repository references it. If export 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(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; + } + + /** + * Ensures that the content for the given claim is accessible outside of this Stateless Process Group by making it available in the backing (on-disk) Content + * Repository, and returns a Content Claim that references it there. 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 #commitExportForExternalUse(ContentClaim)} is called. For any claim that does not belong to this repository, the claim is returned unchanged. + * + * @param claim the claim to make externally accessible + * @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 + */ + public ContentClaim exportForExternalUse(final ContentClaim claim) throws IOException { + if (!(claim instanceof final SpillableContentClaim spillableClaim)) { + 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 #exportForExternalUse(ContentClaim)} + */ + public void commitExportForExternalUse(final ContentClaim claim) { + if (!(claim instanceof final SpillableContentClaim spillableClaim)) { + return; + } + + 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 final byte[] singleByte = new byte[1]; + private UnsynchronizedByteArrayOutputStream buffer = UnsynchronizedByteArrayOutputStream.builder().get(); + private long reserved = 0L; + 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 b) throws IOException { + singleByte[0] = (byte) b; + write(singleByte, 0, 1); + } + + @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 (reserveMemory(len)) { + buffer.write(b, off, len); + reserved += len; + return; + } + + spillOver(); + spillStream.write(b, off, len); + } + + private boolean reserveMemory(final int bytes) { + while (true) { + final long currentMemoryUsed = memoryUsed.get(); + if (currentMemoryUsed > memoryThresholdBytes - bytes) { + return false; + } + + if (memoryUsed.compareAndSet(currentMemoryUsed, currentMemoryUsed + bytes)) { + 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; + 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); + reserved = 0L; + } + + @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(); + } + } + } + + private static final class SpillableContentClaim implements ContentClaim { + private final SpillableResourceClaim resourceClaim; + + private SpillableContentClaim(final boolean lossTolerant) { + this.resourceClaim = new SpillableResourceClaim(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 UnsynchronizedByteArrayOutputStream 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 implements ResourceClaim { + private static final AtomicLong idCounter = new AtomicLong(0L); + + private final String id = String.valueOf(idCounter.getAndIncrement()); + private final boolean lossTolerant; + private volatile UnsynchronizedByteArrayOutputStream 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 boolean lossTolerant) { + this.lossTolerant = lossTolerant; + } + + @Override + public String getId() { + return id; + } + + @Override + public String getContainer() { + return "in-memory"; + } + + @Override + public String getSection() { + return "in-memory"; + } + + @Override + public boolean isLossTolerant() { + return lossTolerant; + } + + @Override + public boolean isWritable() { + return !discarded && (!writeStarted || !writeFinished); + } + + @Override + public boolean isInUse() { + return true; + } + + private long getLength() { + final ContentClaim currentBackingClaim = backingClaim; + if (currentBackingClaim != null) { + return currentBackingClaim.getLength(); + } + + final UnsynchronizedByteArrayOutputStream currentContents = contents; + return currentContents == null ? 0L : currentContents.size(); + } + + private synchronized boolean beginWrite() { + if (writeStarted || discarded || contents != null || backingClaim != null) { + return false; + } + + writeStarted = true; + return true; + } + + private synchronized void finishWrite() { + writeFinished = true; + } + + private boolean isWriteInProgress() { + return writeStarted && !writeFinished; + } + + private synchronized void storeInMemory(final UnsynchronizedByteArrayOutputStream 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 UnsynchronizedByteArrayOutputStream currentContents = contents; + contents = null; + discarded = true; + return currentContents == null ? 0L : currentContents.size(); + } + + private boolean hasInMemoryContents() { + return contents != null; + } + + private void writeInMemoryTo(final OutputStream out) throws IOException { + final UnsynchronizedByteArrayOutputStream currentContents = contents; + if (currentContents != null) { + currentContents.writeTo(out); + } + } + + private InputStream readInMemory() { + final UnsynchronizedByteArrayOutputStream currentContents = contents; + if (currentContents == null) { + return InputStream.nullInputStream(); + } + + return currentContents.toInputStream(); + } + + private synchronized ReleasedContent releaseForCleanup() { + final UnsynchronizedByteArrayOutputStream currentContents = contents; + contents = null; + discarded = true; + + final ContentClaim cleanupClaim; + if (backingClaimOwnership == BackingClaimOwnership.REPOSITORY) { + cleanupClaim = backingClaim; + backingClaimOwnership = BackingClaimOwnership.NONE; + } else { + cleanupClaim = null; + } + + return new ReleasedContent(currentContents == null ? 0L : currentContents.size(), cleanupClaim); + } + + @Override + public boolean equals(final Object o) { + if (this == o) { + return true; + } + + if (o == null || getClass() != o.getClass()) { + return false; + } + + final SpillableResourceClaim that = (SpillableResourceClaim) o; + return Objects.equals(id, that.id); + } + + @Override + public int hashCode() { + return id.hashCode(); + } + } + + 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..182453ce4a04 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 @@ -28,6 +28,7 @@ import org.apache.nifi.controller.metrics.ComponentMetricContext; 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; @@ -35,6 +36,7 @@ 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.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; @@ -332,9 +337,13 @@ private void completeInvocations(final List invocations, final Prove try { updateFlowFileRepository(); } catch (final Exception e) { + rollbackClaimantCounts(); throw new IOException("Failed to update FlowFile Repository after triggering " + this, e); } + commitContentExports(); + incrementedOutputClaims.clear(); + updateProvenanceRepository(statelessProvRepo, event -> true); // Acknowledge the invocations so that the sessions can be committed @@ -352,6 +361,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 { @@ -415,6 +426,7 @@ private void failInvocation(final Invocation invocation, final ProvenanceEventRe try { updateFlowFileRepository(); } catch (final Exception e) { + rollbackClaimantCounts(); throw new IOException("Failed to update FlowFile Repository after triggering " + this, e); } @@ -469,7 +481,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 +489,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 accessible after stateless content is purged. + final FlowFileRecord externallyAccessibleFlowFile = exportContentForExternalUse(outputFlowFile); + + final FlowFileCloneResult cloneResult = ConnectionUtils.clone(externallyAccessibleFlowFile, 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(externallyAccessibleFlowFile, repoRecords, outputPort).ifPresent(cloneProvenanceEvents::add); } } } + private FlowFileRecord exportContentForExternalUse(final FlowFileRecord flowFile) throws IOException { + final ContentClaim contentClaim = flowFile.getContentClaim(); + if (contentClaim == null || !(nifiContentRepository instanceof final DeferredStatelessContentRepository deferredContentRepository)) { + return flowFile; + } + + final ContentClaim externallyAccessibleClaim = deferredContentRepository.exportForExternalUse(contentClaim); + if (externallyAccessibleClaim == contentClaim) { + return flowFile; + } + + preparedContentClaims.add(contentClaim); + return new StandardFlowFileRecord.Builder() + .fromFlowFile(flowFile) + .contentClaim(externallyAccessibleClaim) + .build(); + } + + private void commitContentExports() { + if (!(nifiContentRepository instanceof final DeferredStatelessContentRepository deferredContentRepository)) { + return; + } + + for (final ContentClaim preparedContentClaim : preparedContentClaims) { + deferredContentRepository.commitExportForExternalUse(preparedContentClaim); + } + + preparedContentClaims.clear(); + } + void updateProvenanceRepository(final ProvenanceEventRepository statelessRepo, final Predicate eventFilter) { long firstProvEventId = 0; @@ -539,11 +583,21 @@ 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 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/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..f671cb2c3784 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/repository/SpillableContentRepositoryTest.java @@ -0,0 +1,742 @@ +/* + * 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.StandardContentClaim; +import org.apache.nifi.controller.repository.claim.StandardResourceClaim; +import org.apache.nifi.controller.repository.claim.StandardResourceClaimManager; +import org.apache.nifi.events.EventReporter; +import org.apache.nifi.groups.ProcessGroup; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.mockito.ArgumentCaptor; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +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 java.util.concurrent.atomic.AtomicLong; + +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.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; +import static org.mockito.Mockito.when; + +class SpillableContentRepositoryTest { + private static final long BUDGET = 100L; + + private ResourceClaimManager resourceClaimManager; + private InMemoryContentRepository backingRepository; + private FlowFileRepository flowFileRepository; + private SpillableContentRepository repository; + + @BeforeEach + void setup() throws IOException { + resourceClaimManager = new StandardResourceClaimManager(); + backingRepository = new InMemoryContentRepository(resourceClaimManager); + flowFileRepository = mock(FlowFileRepository.class); + repository = new SpillableContentRepository(backingRepository, flowFileRepository, BUDGET); + repository.initialize(new StandardContentRepositoryContext(resourceClaimManager, EventReporter.NO_OP)); + } + + @Test + void testContentUnderBudgetStaysInMemory() throws IOException { + final byte[] content = bytes(50); + final ContentClaim claim = repository.create(false); + write(claim, content); + + assertTrue(repository.isHeldInMemory(claim)); + assertFalse(repository.isSpilled(claim)); + assertFalse(repository.isAccessible(claim)); + assertEquals(50L, repository.size(claim)); + assertEquals(50L, repository.getInMemoryByteCount()); + assertArrayEquals(content, readAll(claim)); + } + + @Test + void testInMemoryContentSupportsRepeatedReads() throws IOException { + final SpillableContentRepository chunkedRepository = new SpillableContentRepository(backingRepository, flowFileRepository, 4096L); + chunkedRepository.initialize(new StandardContentRepositoryContext(resourceClaimManager, EventReporter.NO_OP)); + final byte[] content = bytes(3000); + final ContentClaim claim = chunkedRepository.create(false); + try (final OutputStream out = chunkedRepository.write(claim)) { + out.write(content); + } + + assertArrayEquals(content, readAll(chunkedRepository, claim)); + assertArrayEquals(content, readAll(chunkedRepository, claim)); + assertEquals(content.length, chunkedRepository.getInMemoryByteCount()); + } + + @Test + void testSingleWriteExceedingBudgetSpills() throws IOException { + final byte[] content = bytes(150); + final ContentClaim claim = repository.create(false); + write(claim, content); + + assertTrue(repository.isSpilled(claim)); + assertFalse(repository.isHeldInMemory(claim)); + assertTrue(repository.isAccessible(claim)); + assertEquals(150L, repository.size(claim)); + assertEquals(0L, repository.getInMemoryByteCount()); + assertArrayEquals(content, readAll(claim)); + } + + @Test + void testMultipleWritesCrossingBudgetSpillMidStream() throws IOException { + final byte[] first = bytes(60); + final byte[] second = bytes(60); + final ContentClaim claim = repository.create(false); + try (final OutputStream out = repository.write(claim)) { + out.write(first); + out.write(second); + } + + assertTrue(repository.isSpilled(claim)); + assertEquals(0L, repository.getInMemoryByteCount()); + assertEquals(120L, repository.size(claim)); + + final byte[] expected = new byte[120]; + System.arraycopy(first, 0, expected, 0, 60); + System.arraycopy(second, 0, expected, 60, 60); + assertArrayEquals(expected, readAll(claim)); + } + + @Test + void testRemoveInMemoryClaimFreesMemory() throws IOException { + final ContentClaim claim = repository.create(false); + write(claim, bytes(50)); + assertEquals(50L, repository.getInMemoryByteCount()); + + repository.remove(claim); + assertEquals(0L, repository.getInMemoryByteCount()); + } + + @Test + void testClaimantCounts() throws IOException { + final ContentClaim claim = repository.create(false); + assertEquals(1, repository.getClaimantCount(claim)); + + assertEquals(2, repository.incrementClaimaintCount(claim)); + assertEquals(1, repository.decrementClaimantCount(claim)); + assertEquals(0, repository.decrementClaimantCount(claim)); + } + + @Test + void testPurgePreservesReferencedClaims() throws IOException { + final byte[] inMemoryContent = bytes(50); + final ContentClaim inMemoryClaim = repository.create(false); + write(inMemoryClaim, inMemoryContent); + + final byte[] spilledContent = bytes(150); + final ContentClaim spilledClaim = repository.create(false); + write(spilledClaim, spilledContent); + + repository.purge(); + + assertEquals(50L, repository.getInMemoryByteCount()); + assertArrayEquals(inMemoryContent, readAll(inMemoryClaim)); + assertArrayEquals(spilledContent, readAll(spilledClaim)); + verify(flowFileRepository, never()).updateRepository(anyList()); + } + + @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)); + + final SpillableContentRepository concurrentRepository = new SpillableContentRepository( + new InMemoryContentRepository(blockingResourceClaimManager), flowFileRepository, BUDGET); + concurrentRepository.initialize(new StandardContentRepositoryContext(blockingResourceClaimManager, EventReporter.NO_OP)); + + 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 testPurgeDiscardsUnreferencedInMemoryClaim() throws IOException { + final ContentClaim claim = repository.create(false); + write(claim, bytes(50)); + repository.decrementClaimantCount(claim); + + repository.purge(); + + assertEquals(0L, repository.getInMemoryByteCount()); + } + + @Test + void testRemoveAfterPurgeDoesNotMakeMemoryCountNegative() throws IOException { + final ContentClaim claim = repository.create(false); + write(claim, bytes(50)); + repository.decrementClaimantCount(claim); + repository.purge(); + + repository.remove(claim); + + assertEquals(0L, repository.getInMemoryByteCount()); + } + + @Test + void testPurgePreservesClaimWithOpenWriter() throws IOException { + final byte[] content = bytes(50); + final ContentClaim claim = repository.create(false); + try (final OutputStream out = repository.write(claim)) { + out.write(content); + repository.decrementClaimantCount(claim); + repository.purge(); + assertEquals(content.length, repository.getInMemoryByteCount()); + } + + assertArrayEquals(content, readAll(claim)); + repository.purge(); + assertEquals(0L, repository.getInMemoryByteCount()); + } + + @Test + void testExportInMemoryClaim() throws IOException { + final byte[] content = bytes(50); + final ContentClaim claim = repository.create(false); + write(claim, content); + + final ContentClaim exported = repository.exportForExternalUse(claim); + assertNotSame(claim, exported); + assertArrayEquals(content, readAll(exported)); + assertEquals(0L, repository.getInMemoryByteCount()); + + assertEquals(1, resourceClaimManager.getClaimantCount(exported.getResourceClaim())); + repository.commitExportForExternalUse(claim); + assertEquals(0, resourceClaimManager.getClaimantCount(exported.getResourceClaim())); + } + + @Test + void testExportInMemoryClaimIsIdempotent() throws IOException { + final byte[] content = bytes(50); + final ContentClaim claim = repository.create(false); + write(claim, content); + + final ContentClaim firstExport = repository.exportForExternalUse(claim); + final ContentClaim secondExport = repository.exportForExternalUse(claim); + + assertSame(firstExport, secondExport); + assertArrayEquals(content, readAll(secondExport)); + assertEquals(1, resourceClaimManager.getClaimantCount(firstExport.getResourceClaim())); + + repository.commitExportForExternalUse(claim); + repository.commitExportForExternalUse(claim); + assertEquals(0, resourceClaimManager.getClaimantCount(firstExport.getResourceClaim())); + } + + @Test + void testExportEmptyInMemoryClaimClosesBackingStream() throws IOException { + final ContentClaim claim = repository.create(false); + write(claim, new byte[0]); + + final ContentClaim exported = repository.exportForExternalUse(claim); + + assertEquals(1, backingRepository.getWriteCount()); + assertEquals(1, backingRepository.getClosedStreamCount()); + assertEquals(0L, exported.getLength()); + repository.commitExportForExternalUse(claim); + } + + @Test + void testExportFailureSubmitsBackingClaimForCleanup() throws IOException { + final ContentClaim claim = repository.create(false); + write(claim, bytes(50)); + backingRepository.failNextWrite(); + + assertThrows(IOException.class, () -> repository.exportForExternalUse(claim)); + + repository.purge(); + verify(flowFileRepository).updateRepository(anyList()); + } + + @Test + void testPurgeCleansExportWhenExternalUpdateFails() throws IOException { + final ContentClaim claim = repository.create(false); + write(claim, bytes(50)); + + final ContentClaim exportedClaim = repository.exportForExternalUse(claim); + repository.incrementClaimaintCount(exportedClaim); + repository.decrementClaimantCount(exportedClaim); + repository.decrementClaimantCount(claim); + repository.purge(); + + final ArgumentCaptor> recordsCaptor = captureRepositoryRecords(); + verify(flowFileRepository).updateRepository(recordsCaptor.capture()); + assertEquals(List.of(exportedClaim), recordsCaptor.getValue().getFirst().getTransientClaims()); + assertEquals(0, resourceClaimManager.getClaimantCount(exportedClaim.getResourceClaim())); + } + + @Test + void testExportSpilledClaimHandsOffBackingClaim() throws IOException { + final byte[] content = bytes(150); + final ContentClaim claim = repository.create(false); + write(claim, content); + + final ContentClaim exported = repository.exportForExternalUse(claim); + assertArrayEquals(content, readAll(exported)); + + assertEquals(1, resourceClaimManager.getClaimantCount(exported.getResourceClaim())); + repository.commitExportForExternalUse(claim); + assertEquals(0, resourceClaimManager.getClaimantCount(exported.getResourceClaim())); + + repository.purge(); + verify(flowFileRepository, never()).updateRepository(anyList()); + } + + @Test + void testPurgeSubmitsUnexportedSpilledClaims() throws IOException { + final ContentClaim claim = repository.create(false); + write(claim, bytes(150)); + final ContentClaim unrelatedClaim = repository.create(false); + final ContentClaim backingClaim = repository.exportForExternalUse(unrelatedClaim); + repository.commitExportForExternalUse(unrelatedClaim); + repository.decrementClaimantCount(claim); + + assertTrue(repository.isSpilled(claim)); + + repository.purge(); + + assertEquals(0L, repository.getInMemoryByteCount()); + + final ArgumentCaptor> captor = captureRepositoryRecords(); + verify(flowFileRepository).updateRepository(captor.capture()); + final List records = captor.getValue(); + assertEquals(1, records.size()); + final RepositoryRecord record = records.get(0); + assertEquals(RepositoryRecordType.CLEANUP_TRANSIENT_CLAIMS, record.getType()); + assertEquals(1, record.getTransientClaims().size()); + + assertFalse(record.getTransientClaims().contains(backingClaim)); + } + + @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 testPurgeCleansOnlyUnreferencedSpilledClaim() throws IOException { + final ContentClaim unreferencedClaim = repository.create(false); + write(unreferencedClaim, bytes(150)); + repository.decrementClaimantCount(unreferencedClaim); + + final byte[] referencedContent = bytes(160); + final ContentClaim referencedClaim = repository.create(false); + write(referencedClaim, referencedContent); + + repository.purge(); + + final ArgumentCaptor> recordsCaptor = ArgumentCaptor.forClass(List.class); + verify(flowFileRepository).updateRepository(recordsCaptor.capture()); + assertEquals(1, recordsCaptor.getValue().getFirst().getTransientClaims().size()); + assertArrayEquals(referencedContent, readAll(referencedClaim)); + } + + @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(); + + assertThrows(IOException.class, () -> out.write(bytes(60))); + out.close(); + + assertEquals(60L, repository.getInMemoryByteCount()); + assertArrayEquals(content, readAll(claim)); + } + + @Test + void testSpilledResourceClaimCanBeRead() throws IOException { + final byte[] content = bytes(150); + final ContentClaim claim = repository.create(false); + write(claim, content); + + assertEquals(content.length, repository.size(claim.getResourceClaim())); + try (final InputStream in = repository.read(claim.getResourceClaim())) { + assertArrayEquals(content, in.readAllBytes()); + } + } + + @Test + void testExportOverwritesDestination(@TempDir final Path tempDirectory) throws IOException { + final byte[] content = bytes(50); + final ContentClaim claim = repository.create(false); + write(claim, content); + final Path destination = tempDirectory.resolve("content"); + Files.write(destination, bytes(100)); + + repository.exportTo(claim, destination, false); + + assertArrayEquals(content, Files.readAllBytes(destination)); + } + + @Test + void testClaimAllowsOnlyOneWriter() throws IOException { + final ContentClaim claim = repository.create(false); + final OutputStream out = repository.write(claim); + + assertThrows(IllegalStateException.class, () -> repository.write(claim)); + out.close(); + assertThrows(IllegalStateException.class, () -> repository.write(claim)); + } + + @Test + void testDeferredRepositoryResolvesChangedThreshold() throws IOException { + final ProcessGroup processGroup = mock(ProcessGroup.class); + final AtomicLong threshold = new AtomicLong(BUDGET); + when(processGroup.resolveStatelessContentMaxHeap()).thenAnswer(invocation -> threshold.get()); + final DeferredStatelessContentRepository deferredRepository = new DeferredStatelessContentRepository( + processGroup, backingRepository, flowFileRepository, resourceClaimManager, EventReporter.NO_OP); + + final ContentClaim inMemoryClaim = deferredRepository.create(false); + assertEquals("in-memory", inMemoryClaim.getResourceClaim().getContainer()); + deferredRepository.decrementClaimantCount(inMemoryClaim); + deferredRepository.purge(); + + threshold.set(0L); + final ContentClaim backingClaim = deferredRepository.create(false); + assertEquals("container", backingClaim.getResourceClaim().getContainer()); + } + + @Test + void testCloneAcrossStates() throws IOException { + final byte[] smallContent = bytes(40); + final ContentClaim small = repository.create(false); + write(small, smallContent); + final ContentClaim smallClone = repository.clone(small, false); + assertArrayEquals(smallContent, readAll(smallClone)); + assertTrue(repository.isHeldInMemory(smallClone)); + + final byte[] largeContent = bytes(150); + final ContentClaim large = repository.create(false); + write(large, largeContent); + final ContentClaim largeClone = repository.clone(large, false); + assertArrayEquals(largeContent, readAll(largeClone)); + assertTrue(repository.isSpilled(largeClone)); + } + + @Test + void testImportFromAndExportTo() throws IOException { + final byte[] content = bytes(150); + final ContentClaim claim = repository.create(false); + final long imported = repository.importFrom(new ByteArrayInputStream(content), claim); + assertEquals(150L, imported); + + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + repository.exportTo(claim, out); + assertArrayEquals(content, out.toByteArray()); + } + + 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 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; + } + + /** + * A minimal on-heap {@link ContentRepository} used to stand in for the NiFi Content Repository that a {@link SpillableContentRepository} spills to. It shares + * the {@link ResourceClaimManager} used by the repository under test so that claimant counts can be asserted uniformly. + */ + private static final class InMemoryContentRepository implements ContentRepository { + private final ResourceClaimManager resourceClaimManager; + private final Map contents = new HashMap<>(); + private final AtomicInteger idGenerator = new AtomicInteger(0); + private final AtomicInteger writeCount = new AtomicInteger(0); + private final AtomicInteger closedStreamCount = new AtomicInteger(0); + private boolean failNextWrite; + + private InMemoryContentRepository(final ResourceClaimManager resourceClaimManager) { + this.resourceClaimManager = resourceClaimManager; + } + + @Override + public void initialize(final ContentRepositoryContext context) { + } + + @Override + public void shutdown() { + } + + @Override + public Set getContainerNames() { + return Set.of("container"); + } + + @Override + public long getContainerCapacity(final String containerName) { + return 0; + } + + @Override + public long getContainerUsableSpace(final String containerName) { + return 0; + } + + @Override + public String getContainerFileStoreName(final String containerName) { + return "container"; + } + + @Override + public ContentClaim create(final boolean lossTolerant) { + final ResourceClaim resourceClaim = new StandardResourceClaim(resourceClaimManager, "container", "section", "backing-" + idGenerator.getAndIncrement(), lossTolerant); + final StandardContentClaim contentClaim = new StandardContentClaim(resourceClaim, 0L); + contentClaim.setLength(0L); + resourceClaimManager.incrementClaimantCount(resourceClaim); + return contentClaim; + } + + @Override + public int incrementClaimaintCount(final ContentClaim claim) { + return resourceClaimManager.incrementClaimantCount(claim.getResourceClaim()); + } + + @Override + public int getClaimantCount(final ContentClaim claim) { + return resourceClaimManager.getClaimantCount(claim.getResourceClaim()); + } + + @Override + public int decrementClaimantCount(final ContentClaim claim) { + return resourceClaimManager.decrementClaimantCount(claim.getResourceClaim()); + } + + @Override + public boolean remove(final ContentClaim claim) { + contents.remove(claim.getResourceClaim()); + 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); + } + + return clone; + } + + @Override + public long importFrom(final Path content, final ContentClaim claim) { + throw new UnsupportedOperationException(); + } + + @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) { + throw new UnsupportedOperationException(); + } + + @Override + public long exportTo(final ContentClaim claim, final Path destination, final boolean append, final long offset, final long length) { + throw new UnsupportedOperationException(); + } + + @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) { + throw new UnsupportedOperationException(); + } + + @Override + public long size(final ContentClaim claim) { + return claim.getLength(); + } + + @Override + public long size(final ResourceClaim claim) { + final byte[] data = contents.get(claim); + return data == null ? 0 : data.length; + } + + @Override + public InputStream read(final ContentClaim claim) { + final byte[] data = contents.getOrDefault(claim.getResourceClaim(), new byte[0]); + return new ByteArrayInputStream(data); + } + + @Override + public InputStream read(final ResourceClaim claim) { + return new ByteArrayInputStream(contents.getOrDefault(claim, new byte[0])); + } + + @Override + public OutputStream write(final ContentClaim claim) { + final StandardContentClaim standardContentClaim = (StandardContentClaim) claim; + writeCount.incrementAndGet(); + if (failNextWrite) { + failNextWrite = false; + return new OutputStream() { + @Override + public void write(final int value) throws IOException { + throw new IOException("Write failed"); + } + + @Override + public void write(final byte[] bytes, final int offset, final int length) throws IOException { + throw new IOException("Write failed"); + } + + @Override + public void close() { + closedStreamCount.incrementAndGet(); + } + }; + } + + return new ByteArrayOutputStream() { + @Override + public void close() { + final byte[] data = toByteArray(); + contents.put(standardContentClaim.getResourceClaim(), data); + standardContentClaim.setLength(data.length); + closedStreamCount.incrementAndGet(); + } + }; + } + + private void failNextWrite() { + failNextWrite = true; + } + + private int getWriteCount() { + return writeCount.get(); + } + + private int getClosedStreamCount() { + return closedStreamCount.get(); + } + + @Override + public void purge() { + contents.clear(); + } + + @Override + public void cleanup() { + } + + @Override + public boolean isAccessible(final ContentClaim contentClaim) { + return contentClaim != null && contents.containsKey(contentClaim.getResourceClaim()); + } + } +} 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..c3301a280e24 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 0; + } + + @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..cd2a20d52325 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,16 +75,22 @@ 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.times; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; public class TestStatelessFlowTask { @@ -94,6 +107,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 +148,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 +176,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 +225,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 +251,100 @@ public void testCreateOutputRecordsOnSuccessWithOneOutput() { } @Test - public void testCreateOutputRecordsOnSuccessWithOnePortTwoConnections() { + public void testCreateOutputRecordsExportsSharedInMemoryClaimOnce() throws IOException { + final byte[] content = "shared content".getBytes(StandardCharsets.UTF_8); + final ByteArrayContentRepository backingRepository = new ByteArrayContentRepository(); + backingRepository.initialize(new StandardContentRepositoryContext(resourceClaimManager, EventReporter.NO_OP)); + final ProcessGroup processGroup = mock(ProcessGroup.class); + when(processGroup.resolveStatelessContentMaxHeap()).thenReturn(1024L); + final DeferredStatelessContentRepository deferredRepository = new DeferredStatelessContentRepository( + processGroup, backingRepository, flowFileRepository, resourceClaimManager, EventReporter.NO_OP); + + 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 ByteArrayContentRepository backingRepository = new ByteArrayContentRepository(); + backingRepository.initialize(new StandardContentRepositoryContext(resourceClaimManager, EventReporter.NO_OP)); + final ProcessGroup processGroup = mock(ProcessGroup.class); + when(processGroup.resolveStatelessContentMaxHeap()).thenReturn(1024L); + final DeferredStatelessContentRepository deferredRepository = new DeferredStatelessContentRepository( + processGroup, backingRepository, flowFileRepository, resourceClaimManager, EventReporter.NO_OP); + + 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())); + verify(flowFileRepository, times(2)).updateRepository(any()); + } + + @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 +388,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 +454,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 +470,7 @@ public void testUpdateClaimantCountSingleOutput() { } @Test - public void testUpdateClaimantCountWithClone() { + public void testUpdateClaimantCountWithClone() throws IOException { final Connection conn1 = mockConnection(); final Connection conn2 = mockConnection(); @@ -389,7 +491,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 +514,18 @@ private Connection mockConnection() { return connection; } + 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..689f509570f2 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("Illegal value proposed for Max In-Memory Heap Percentage: " + heapPercentage); + } + } catch (final NumberFormatException e) { + throw new IllegalArgumentException("Illegal value proposed for Max In-Memory Heap Percentage: " + 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 46e8ab35d7ad..88c63d9e2cf6 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 @@ -2836,6 +2836,9 @@ private ProcessGroupDTO createConciseProcessGroupDto(final ProcessGroup group) { dto.setExecutionEngine(group.getExecutionEngine().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) { @@ -4827,6 +4830,8 @@ public ProcessGroupDTO copy(final ProcessGroupDTO original, final boolean deep) copy.setExecutionEngine(original.getExecutionEngine()); 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..870822b47681 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..f64600bc5fd1 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,82 @@ 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).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 +296,8 @@ describe('EditProcessGroup', () => { }, executionEngine: 'INHERITED', maxConcurrentTasks: 1, - statelessFlowTimeout: '1 min' + statelessFlowTimeout: '1 min', + statelessFlowFileContentInMemoryMax: '0 B' } } }; @@ -303,7 +382,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..4f61f5c4ea69 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 ?? 0; 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..9d2e72f2eb76 --- /dev/null +++ b/nifi-stateless/nifi-stateless-bundle/nifi-stateless-engine/src/test/java/org/apache/nifi/stateless/engine/StandardExecutionProgressTest.java @@ -0,0 +1,145 @@ +/* + * 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 testOutputClaimReleasedWhenAborted() { + final TestContext context = createTestContext(); + context.progress().enqueueTriggerResult(() -> { }, failure -> { }); + final TriggerResult result = context.results().remove(); + + result.abort(new IOException("Processing failed")); + + verify(context.contentRepository()).decrementClaimantCount(context.contentClaim()); + verify(context.contentRepository()).purge(); + verify(context.purgeAction()).purge(); + } + + @Test + void testEachAbortedResultPurgesReleasedClaims() { + final TestContext context = createTestContext(); + context.progress().enqueueTriggerResult(() -> { }, failure -> { }); + context.progress().enqueueTriggerResult(() -> { }, failure -> { }); + final TriggerResult firstResult = context.results().remove(); + final TriggerResult secondResult = context.results().remove(); + + firstResult.abort(new IOException("First processing failure")); + secondResult.abort(new IOException("Second processing failure")); + + verify(context.contentRepository(), times(2)).decrementClaimantCount(context.contentClaim()); + verify(context.contentRepository(), times(2)).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-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..8775599aa9a1 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 @@ -70,6 +70,7 @@ import org.apache.nifi.web.api.dto.RemoteProcessGroupDTO; import org.apache.nifi.web.api.dto.ReportingTaskDTO; import org.apache.nifi.web.api.dto.RevisionDTO; +import org.apache.nifi.web.api.dto.SnippetDTO; import org.apache.nifi.web.api.dto.VerifyConfigRequestDTO; import org.apache.nifi.web.api.dto.VerifyConnectorConfigStepRequestDTO; import org.apache.nifi.web.api.dto.VersionControlInformationDTO; @@ -135,6 +136,7 @@ import org.apache.nifi.web.api.entity.ReportingTaskRunStatusEntity; import org.apache.nifi.web.api.entity.ReportingTasksEntity; import org.apache.nifi.web.api.entity.ScheduleComponentsEntity; +import org.apache.nifi.web.api.entity.SnippetEntity; import org.apache.nifi.web.api.entity.StartVersionControlRequestEntity; import org.apache.nifi.web.api.entity.VerifyConfigRequestEntity; import org.apache.nifi.web.api.entity.VerifyConnectorConfigStepRequestEntity; @@ -155,6 +157,7 @@ import java.util.Collection; import java.util.Collections; import java.util.Date; +import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; import java.util.List; @@ -2890,4 +2893,38 @@ 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); + } + + public SnippetEntity moveProcessGroup(final ProcessGroupEntity groupToMove, final String destinationGroupId) throws NiFiClientException, IOException { + final Map processGroupRevisions = new HashMap<>(); + processGroupRevisions.put(groupToMove.getId(), groupToMove.getRevision()); + + final SnippetDTO snippetDto = new SnippetDTO(); + snippetDto.setParentGroupId(groupToMove.getComponent().getParentGroupId()); + snippetDto.setProcessGroups(processGroupRevisions); + + final SnippetEntity snippet = new SnippetEntity(); + snippet.setSnippet(snippetDto); + final SnippetEntity createdSnippet = nifiClient.getSnippetClient().createSnippet(snippet); + + createdSnippet.getSnippet().setParentGroupId(destinationGroupId); + return nifiClient.getSnippetClient().updateSnippet(createdSnippet); + } } 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..418c0aee9f5a --- /dev/null +++ b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/stateless/StatelessInMemoryContentIT.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.tests.system.stateless; + +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 java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Map; +import java.util.stream.Stream; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +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"; + + @Override + protected boolean isAllowFactoryReuse() { + return false; + } + + @Override + protected boolean isDestroyEnvironmentAfterEachTest() { + return true; + } + + @Test + public void testContentUnderBudgetStaysInMemory() throws NiFiClientException, IOException, InterruptedException { + final SelfContainedFlow flow = createSelfContainedTransformFlow(LARGE_BUDGET); + final long contentBytesBeforeStart = contentBytesOnDisk(); + + getClientUtil().startProcessGroupComponents(flow.groupId()); + + waitFor(() -> Files.exists(flow.markerFile())); + assertEquals(contentBytesBeforeStart, contentBytesOnDisk()); + 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 long contentBytesBeforeStart = contentBytesOnDisk(); + + getClientUtil().startProcessGroupComponents(flow.groupId()); + + waitFor(() -> Files.exists(flow.markerFile())); + assertTrue(contentBytesOnDisk() > contentBytesBeforeStart); + 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. + assertThrows(NiFiClientException.class, () -> getClientUtil().setStatelessFlowFileContentInMemoryMax(statelessGroup, TINY_BUDGET)); + + getClientUtil().stopProcessGroupComponents(groupId); + + final ProcessGroupEntity stoppedGroup = getClientUtil().setStatelessFlowFileContentInMemoryMax(statelessGroup, TINY_BUDGET); + assertEquals(TINY_BUDGET, stoppedGroup.getComponent().getStatelessFlowFileContentInMemoryMax()); + + final long contentBytesBeforeRestart = contentBytesOnDisk(); + getClientUtil().startProcessGroupComponents(groupId); + waitFor(() -> getProcessorFlowFilesIn(terminate.getId()) >= 2); + assertTrue(contentBytesOnDisk() > contentBytesBeforeRestart); + 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 on disk. The transformed content is verified byte-for-byte. For the spill case, a Sleep + * processor keeps the transformed FlowFile inside the Stateless group long enough to confirm that its claim is written to disk before it reaches the Output + * Port boundary. + */ + 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); + if (verifySpill) { + final ProcessorEntity sleep = getClientUtil().createProcessor("Sleep", groupId); + getClientUtil().updateProcessorProperties(sleep, Map.of("onTrigger Sleep Time", "10 sec")); + getClientUtil().createConnection(reverse, sleep, SUCCESS, groupId); + getClientUtil().createConnection(sleep, outputPort, SUCCESS); + getClientUtil().waitForValidProcessor(sleep.getId()); + } else { + 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 long contentBytesBeforeStateless = contentBytesOnDisk(); + getClientUtil().startProcessGroupComponents(groupId); + + if (verifySpill) { + waitFor(() -> contentBytesOnDisk() > contentBytesBeforeStateless); + assertEquals(0, getConnectionQueueSize(outputToTerminate.getId())); + } + + 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"); + statelessGroup.getComponent().setMaxConcurrentTasks(4); + 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); + getClientUtil().updateProcessorRunDuration(reverseFirst, 25); + getClientUtil().updateProcessorRunDuration(append, 25); + getClientUtil().updateProcessorRunDuration(reverseSecond, 25); + + 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 ProcessorEntity sleep = getClientUtil().createProcessor("Sleep", groupId); + getClientUtil().updateProcessorProperties(sleep, Map.of("onTrigger Sleep Time", "10 sec")); + 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, sleep, SUCCESS, groupId); + getClientUtil().createConnection(sleep, 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(sleep.getId()); + + return new SelfContainedFlow(groupId, matchedTerminate.getId(), unmatchedTerminate.getId(), markerFile); + } + + private int getProcessorFlowFilesIn(final String processorId) throws NiFiClientException, IOException { + return getNifiClient().getProcessorClient().getProcessor(processorId).getStatus().getAggregateSnapshot().getFlowFilesIn(); + } + + private long contentBytesOnDisk() throws IOException { + final File contentRepository = new File(getNiFiInstance().getInstanceDirectory(), "content_repository"); + if (!contentRepository.exists()) { + return 0L; + } + + try (final Stream paths = Files.walk(contentRepository.toPath())) { + return paths.filter(Files::isRegularFile) + .mapToLong(path -> path.toFile().length()) + .sum(); + } + } + + private record SelfContainedFlow(String groupId, String matchedTerminateId, String unmatchedTerminateId, Path markerFile) { + } +} diff --git a/pom.xml b/pom.xml index 84379ce49a3f..3c154b948374 100644 --- a/pom.xml +++ b/pom.xml @@ -118,7 +118,7 @@ v24.14.1 - 2.12.0 + 2.12.0-SNAPSHOT 2.4.0