This is an automated email from the ASF dual-hosted git repository. joewitt pushed a commit to branch NIFI-16322 in repository https://gitbox.apache.org/repos/asf/nifi.git
commit 40752b9c6292996cd679cfa410bd864163da6d7c Author: Joseph Witt <[email protected]> AuthorDate: Tue Sep 8 14:17:55 2026 -0700 NIFI-16322 Stream GenerateFlowFile content instead of a full-size byte array Co-authored-by: Cursor <[email protected]> --- .../nifi/processors/standard/GenerateFlowFile.java | 131 +++++++++++++-------- .../processors/standard/TestGenerateFlowFile.java | 73 +++++++++++- 2 files changed, 152 insertions(+), 52 deletions(-) diff --git a/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/GenerateFlowFile.java b/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/GenerateFlowFile.java index 3a46949cbd5..f173a97df5b 100644 --- a/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/GenerateFlowFile.java +++ b/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/GenerateFlowFile.java @@ -42,6 +42,8 @@ import org.apache.nifi.processor.Relationship; import org.apache.nifi.processor.util.StandardValidators; import org.apache.nifi.scheduling.SchedulingStrategy; +import java.io.IOException; +import java.io.OutputStream; import java.nio.charset.Charset; import java.util.ArrayList; import java.util.Collection; @@ -67,7 +69,9 @@ import java.util.concurrent.atomic.AtomicReference; @DefaultSchedule(strategy = SchedulingStrategy.TIMER_DRIVEN, period = "1 min") public class GenerateFlowFile extends AbstractProcessor { - private final AtomicReference<byte[]> data = new AtomicReference<>(); + private static final int BUFFER_SIZE = 8192; + + private final AtomicReference<GeneratedData> generatedData = new AtomicReference<>(); public static final String DATA_FORMAT_BINARY = "Binary"; public static final String DATA_FORMAT_TEXT = "Text"; @@ -77,7 +81,7 @@ public class GenerateFlowFile extends AbstractProcessor { .description("The size of the file that will be used") .required(true) .defaultValue("0B") - .addValidator(StandardValidators.DATA_SIZE_VALIDATOR) + .addValidator(StandardValidators.createDataSizeBoundsValidator(0, Integer.MAX_VALUE)) .expressionLanguageSupported(ExpressionLanguageScope.ENVIRONMENT) .build(); public static final PropertyDescriptor BATCH_SIZE = new PropertyDescriptor.Builder() @@ -96,8 +100,8 @@ public class GenerateFlowFile extends AbstractProcessor { .build(); public static final PropertyDescriptor UNIQUE_FLOWFILES = new PropertyDescriptor.Builder() .name("Unique FlowFiles") - .description("If true, each FlowFile that is generated will be unique. If false, a random value will be generated and all FlowFiles " - + "will get the same content but this offers much higher throughput") + .description("If true, each FlowFile that is generated will be unique. If false, all generated FlowFiles will have the same content. " + + "When Unique FlowFiles is false, the first FlowFile in a batch is written and additional FlowFiles clone that content.") .required(true) .allowableValues("true", "false") .defaultValue("false") @@ -169,10 +173,10 @@ public class GenerateFlowFile extends AbstractProcessor { @OnScheduled public void onScheduled(final ProcessContext context) { - if (context.getProperty(UNIQUE_FLOWFILES).asBoolean()) { - this.data.set(null); - } else if (!context.getProperty(CUSTOM_TEXT).isSet()) { - this.data.set(generateData(context)); + if (context.getProperty(UNIQUE_FLOWFILES).asBoolean() || context.getProperty(CUSTOM_TEXT).isSet()) { + generatedData.set(null); + } else { + generatedData.set(getGeneratedData(context)); } } @@ -191,45 +195,44 @@ public class GenerateFlowFile extends AbstractProcessor { return results; } - private byte[] generateData(final ProcessContext context) { - final int byteCount = context.getProperty(FILE_SIZE).evaluateAttributeExpressions().asDataSize(DataUnit.B).intValue(); - - final Random random = new Random(); - final byte[] array = new byte[byteCount]; - if (context.getProperty(DATA_FORMAT).getValue().equals(DATA_FORMAT_BINARY)) { - random.nextBytes(array); - } else { - for (int i = 0; i < array.length; i++) { - final int index = random.nextInt(TEXT_CHARS.length); - array[i] = (byte) TEXT_CHARS[index]; - } - } - - return array; + private GeneratedData getGeneratedData(final ProcessContext context) { + final long byteCount = context.getProperty(FILE_SIZE).evaluateAttributeExpressions().asDataSize(DataUnit.B).longValue(); + final boolean binary = context.getProperty(DATA_FORMAT).getValue().equals(DATA_FORMAT_BINARY); + return new GeneratedData(byteCount, binary, new Random().nextLong()); } @Override public void onTrigger(final ProcessContext context, final ProcessSession session) { - final byte[] data; final boolean uniqueData = context.getProperty(UNIQUE_FLOWFILES).asBoolean(); + final int batchSize = context.getProperty(BATCH_SIZE).asInteger(); + final Map<String, String> generatedAttributes = getGeneratedAttributes(context); + if (uniqueData) { - data = new byte[0]; - } else { - if (context.getProperty(CUSTOM_TEXT).isSet()) { - final Charset charset = Charset.forName(context.getProperty(CHARSET).getValue()); - data = context.getProperty(CUSTOM_TEXT).evaluateAttributeExpressions().getValue().getBytes(charset); - } else { - data = this.data.get(); + for (int i = 0; i < batchSize; i++) { + session.transfer(createGeneratedFlowFile(session, getGeneratedData(context), generatedAttributes), SUCCESS); } + return; + } + + final FlowFile first; + if (context.getProperty(CUSTOM_TEXT).isSet()) { + first = createCustomTextFlowFile(context, session, generatedAttributes); + } else { + first = createGeneratedFlowFile(session, generatedData.get(), generatedAttributes); + } + + for (int i = 1; i < batchSize; i++) { + session.transfer(session.clone(first), SUCCESS); } + session.transfer(first, SUCCESS); + } - Map<PropertyDescriptor, String> processorProperties = context.getProperties(); - Map<String, String> generatedAttributes = new HashMap<>(); - for (final Map.Entry<PropertyDescriptor, String> entry : processorProperties.entrySet()) { - PropertyDescriptor property = entry.getKey(); + private Map<String, String> getGeneratedAttributes(final ProcessContext context) { + final Map<String, String> generatedAttributes = new HashMap<>(); + for (final Map.Entry<PropertyDescriptor, String> entry : context.getProperties().entrySet()) { + final PropertyDescriptor property = entry.getKey(); if (property.isDynamic() && property.isExpressionLanguageSupported()) { - String dynamicValue = context.getProperty(property).evaluateAttributeExpressions().getValue(); - generatedAttributes.put(property.getName(), dynamicValue); + generatedAttributes.put(property.getName(), context.getProperty(property).evaluateAttributeExpressions().getValue()); } } @@ -237,17 +240,53 @@ public class GenerateFlowFile extends AbstractProcessor { generatedAttributes.put(CoreAttributes.MIME_TYPE.key(), context.getProperty(MIME_TYPE).getValue()); } - for (int i = 0; i < context.getProperty(BATCH_SIZE).asInteger(); i++) { - FlowFile flowFile = session.create(); - final byte[] writtenData = uniqueData ? generateData(context) : data; - if (writtenData.length > 0) { - flowFile = session.write(flowFile, out -> out.write(writtenData)); - } - flowFile = session.putAllAttributes(flowFile, generatedAttributes); + return generatedAttributes; + } + + private FlowFile createCustomTextFlowFile(final ProcessContext context, final ProcessSession session, final Map<String, String> generatedAttributes) { + final Charset charset = Charset.forName(context.getProperty(CHARSET).getValue()); + final byte[] customData = context.getProperty(CUSTOM_TEXT).evaluateAttributeExpressions().getValue().getBytes(charset); + FlowFile flowFile = session.create(); + if (customData.length > 0) { + flowFile = session.write(flowFile, out -> out.write(customData)); + } + return finishCreatedFlowFile(session, flowFile, generatedAttributes); + } - session.getProvenanceReporter().create(flowFile); - session.transfer(flowFile, SUCCESS); + private FlowFile createGeneratedFlowFile(final ProcessSession session, final GeneratedData data, final Map<String, String> generatedAttributes) { + FlowFile flowFile = session.create(); + if (data.byteCount() > 0) { + flowFile = session.write(flowFile, out -> writeGeneratedData(out, data)); } + return finishCreatedFlowFile(session, flowFile, generatedAttributes); + } + + private FlowFile finishCreatedFlowFile(final ProcessSession session, FlowFile flowFile, final Map<String, String> generatedAttributes) { + flowFile = session.putAllAttributes(flowFile, generatedAttributes); + session.getProvenanceReporter().create(flowFile); + return flowFile; + } + + private static void writeGeneratedData(final OutputStream outputStream, final GeneratedData generatedData) throws IOException { + final Random random = new Random(generatedData.seed()); + final byte[] buffer = new byte[(int) Math.min(BUFFER_SIZE, generatedData.byteCount())]; + long remaining = generatedData.byteCount(); + + while (remaining > 0) { + final int length = (int) Math.min(buffer.length, remaining); + if (generatedData.binary()) { + random.nextBytes(buffer); + } else { + for (int i = 0; i < length; i++) { + buffer[i] = (byte) TEXT_CHARS[random.nextInt(TEXT_CHARS.length)]; + } + } + outputStream.write(buffer, 0, length); + remaining -= length; + } + } + + private record GeneratedData(long byteCount, boolean binary, long seed) { } @Override diff --git a/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestGenerateFlowFile.java b/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestGenerateFlowFile.java index 89e4a82aa9f..2941b8e6140 100644 --- a/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestGenerateFlowFile.java +++ b/nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestGenerateFlowFile.java @@ -23,15 +23,22 @@ import org.apache.nifi.util.TestRunners; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import java.util.Arrays; +import java.util.List; import java.util.Map; +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.assertTrue; /** * Unit tests for the GenerateFlowFile processor. */ public class TestGenerateFlowFile { + private static final String STREAMED_FILE_SIZE = "9 KB"; + private static final long STREAMED_FILE_SIZE_BYTES = 9 * 1024; + private TestRunner runner; @BeforeEach @@ -41,21 +48,23 @@ public class TestGenerateFlowFile { @Test public void testGenerateCustomText() { - runner.setProperty(GenerateFlowFile.FILE_SIZE, "100MB"); + runner.setProperty(GenerateFlowFile.FILE_SIZE, "1B"); runner.setProperty(GenerateFlowFile.DATA_FORMAT, GenerateFlowFile.DATA_FORMAT_TEXT); runner.setProperty(GenerateFlowFile.CUSTOM_TEXT, "This is my custom text!"); + runner.setProperty(GenerateFlowFile.BATCH_SIZE, "2"); + runner.run(); - runner.assertTransferCount(GenerateFlowFile.SUCCESS, 1); - MockFlowFile generatedFlowFile = runner.getFlowFilesForRelationship(GenerateFlowFile.SUCCESS).getFirst(); - generatedFlowFile.assertContentEquals("This is my custom text!"); - generatedFlowFile.assertAttributeNotExists("mime.type"); + runner.assertTransferCount(GenerateFlowFile.SUCCESS, 2); + runner.getFlowFilesForRelationship(GenerateFlowFile.SUCCESS).get(0).assertContentEquals("This is my custom text!"); + runner.getFlowFilesForRelationship(GenerateFlowFile.SUCCESS).get(1).assertContentEquals("This is my custom text!"); + runner.getFlowFilesForRelationship(GenerateFlowFile.SUCCESS).get(0).assertAttributeNotExists("mime.type"); } @Test public void testInvalidCustomText() { - runner.setProperty(GenerateFlowFile.FILE_SIZE, "100MB"); + runner.setProperty(GenerateFlowFile.FILE_SIZE, "1B"); runner.setProperty(GenerateFlowFile.DATA_FORMAT, GenerateFlowFile.DATA_FORMAT_BINARY); runner.setProperty(GenerateFlowFile.CUSTOM_TEXT, "This is my custom text!"); runner.assertNotValid(); @@ -65,6 +74,34 @@ public class TestGenerateFlowFile { runner.assertNotValid(); } + @Test + public void testFileSizeLargerThanIntegerMaxIsInvalid() { + // Configuration only — do not runner.run() with this size. GitHub runners cannot host a 3 GB FlowFile. + runner.setProperty(GenerateFlowFile.FILE_SIZE, "3 GB"); + + runner.assertNotValid(); + } + + @Test + public void testGenerateNonUniqueBinaryContentStreamed() { + assertStreamedContent(false, GenerateFlowFile.DATA_FORMAT_BINARY); + } + + @Test + public void testGenerateNonUniqueTextContentStreamed() { + assertStreamedContent(false, GenerateFlowFile.DATA_FORMAT_TEXT); + } + + @Test + public void testGenerateUniqueBinaryContentStreamed() { + assertStreamedContent(true, GenerateFlowFile.DATA_FORMAT_BINARY); + } + + @Test + public void testGenerateUniqueTextContentStreamed() { + assertStreamedContent(true, GenerateFlowFile.DATA_FORMAT_TEXT); + } + @Test public void testDynamicPropertiesToAttributes() { runner.setProperty(GenerateFlowFile.FILE_SIZE, "1B"); @@ -136,4 +173,28 @@ public class TestGenerateFlowFile { assertEquals(expectedRenamed, propertyMigrationResult.getPropertiesRenamed()); } + private void assertStreamedContent(final boolean unique, final String dataFormat) { + runner.setProperty(GenerateFlowFile.FILE_SIZE, STREAMED_FILE_SIZE); + runner.setProperty(GenerateFlowFile.DATA_FORMAT, dataFormat); + runner.setProperty(GenerateFlowFile.UNIQUE_FLOWFILES, Boolean.toString(unique)); + runner.setProperty(GenerateFlowFile.BATCH_SIZE, unique ? "2" : "3"); + + runner.run(); + + runner.assertAllFlowFilesTransferred(GenerateFlowFile.SUCCESS, unique ? 2 : 3); + final List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(GenerateFlowFile.SUCCESS); + final MockFlowFile first = flowFiles.getFirst(); + for (final MockFlowFile flowFile : flowFiles) { + assertEquals(STREAMED_FILE_SIZE_BYTES, flowFile.getSize()); + } + + if (unique) { + assertFalse(Arrays.equals(first.getData(), flowFiles.get(1).getData())); + } else { + for (int i = 1; i < flowFiles.size(); i++) { + assertArrayEquals(first.getData(), flowFiles.get(i).getData()); + } + } + } + }
