Search in sources :

Example 1 with ACCEPT_AND_CONTINUE

use of org.apache.nifi.processor.FlowFileFilter.FlowFileFilterResult.ACCEPT_AND_CONTINUE in project nifi by apache.

the class Wait method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    final ComponentLog logger = getLogger();
    // Signal id is computed from attribute 'RELEASE_SIGNAL_IDENTIFIER' with expression language support
    final PropertyValue signalIdProperty = context.getProperty(RELEASE_SIGNAL_IDENTIFIER);
    final Integer bufferCount = context.getProperty(WAIT_BUFFER_COUNT).asInteger();
    final Map<Relationship, List<FlowFile>> processedFlowFiles = new HashMap<>();
    final Function<Relationship, List<FlowFile>> getFlowFilesFor = r -> processedFlowFiles.computeIfAbsent(r, k -> new ArrayList<>());
    final AtomicReference<String> targetSignalId = new AtomicReference<>();
    final AtomicInteger bufferedCount = new AtomicInteger(0);
    final List<FlowFile> failedFilteringFlowFiles = new ArrayList<>();
    final Supplier<FlowFileFilter.FlowFileFilterResult> acceptResultSupplier = () -> bufferedCount.incrementAndGet() == bufferCount ? ACCEPT_AND_TERMINATE : ACCEPT_AND_CONTINUE;
    final List<FlowFile> flowFiles = session.get(f -> {
        final String fSignalId = signalIdProperty.evaluateAttributeExpressions(f).getValue();
        // if the computed value is null, or empty, we transfer the FlowFile to failure relationship
        if (StringUtils.isBlank(fSignalId)) {
            // We can't penalize f before getting it from session, so keep it in a temporal list.
            logger.error("FlowFile {} has no attribute for given Release Signal Identifier", new Object[] { f });
            failedFilteringFlowFiles.add(f);
            return ACCEPT_AND_CONTINUE;
        }
        final String targetSignalIdStr = targetSignalId.get();
        if (targetSignalIdStr == null) {
            // This is the first one.
            targetSignalId.set(fSignalId);
            return acceptResultSupplier.get();
        }
        if (targetSignalIdStr.equals(fSignalId)) {
            return acceptResultSupplier.get();
        }
        return REJECT_AND_CONTINUE;
    });
    final String attributeCopyMode = context.getProperty(ATTRIBUTE_COPY_MODE).getValue();
    final boolean replaceOriginalAttributes = ATTRIBUTE_COPY_REPLACE.getValue().equals(attributeCopyMode);
    final AtomicReference<Signal> signalRef = new AtomicReference<>();
    // This map contains original counts before those are consumed to release incoming FlowFiles.
    final HashMap<String, Long> originalSignalCounts = new HashMap<>();
    final Consumer<FlowFile> transferToFailure = flowFile -> {
        flowFile = session.penalize(flowFile);
        getFlowFilesFor.apply(REL_FAILURE).add(flowFile);
    };
    final Consumer<Entry<Relationship, List<FlowFile>>> transferFlowFiles = routedFlowFiles -> {
        Relationship relationship = routedFlowFiles.getKey();
        if (REL_WAIT.equals(relationship)) {
            final String waitMode = context.getProperty(WAIT_MODE).getValue();
            if (WAIT_MODE_KEEP_IN_UPSTREAM.getValue().equals(waitMode)) {
                // Transfer to self.
                relationship = Relationship.SELF;
            }
        }
        final List<FlowFile> flowFilesWithSignalAttributes = routedFlowFiles.getValue().stream().map(f -> copySignalAttributes(session, f, signalRef.get(), originalSignalCounts, replaceOriginalAttributes)).collect(Collectors.toList());
        session.transfer(flowFilesWithSignalAttributes, relationship);
    };
    failedFilteringFlowFiles.forEach(f -> {
        flowFiles.remove(f);
        transferToFailure.accept(f);
    });
    if (flowFiles.isEmpty()) {
        // If there was nothing but failed FlowFiles while filtering, transfer those and end immediately.
        processedFlowFiles.entrySet().forEach(transferFlowFiles);
        return;
    }
    // the cache client used to interact with the distributed cache
    final AtomicDistributedMapCacheClient cache = context.getProperty(DISTRIBUTED_CACHE_SERVICE).asControllerService(AtomicDistributedMapCacheClient.class);
    final WaitNotifyProtocol protocol = new WaitNotifyProtocol(cache);
    final String signalId = targetSignalId.get();
    final Signal signal;
    // get notifying signal
    try {
        signal = protocol.getSignal(signalId);
        if (signal != null) {
            originalSignalCounts.putAll(signal.getCounts());
        }
        signalRef.set(signal);
    } catch (final IOException e) {
        throw new ProcessException(String.format("Failed to get signal for %s due to %s", signalId, e), e);
    }
    String targetCounterName = null;
    long targetCount = 1;
    int releasableFlowFileCount = 1;
    final List<FlowFile> candidates = new ArrayList<>();
    for (FlowFile flowFile : flowFiles) {
        // Set wait start timestamp if it's not set yet
        String waitStartTimestamp = flowFile.getAttribute(WAIT_START_TIMESTAMP);
        if (waitStartTimestamp == null) {
            waitStartTimestamp = String.valueOf(System.currentTimeMillis());
            flowFile = session.putAttribute(flowFile, WAIT_START_TIMESTAMP, waitStartTimestamp);
        }
        long lWaitStartTimestamp;
        try {
            lWaitStartTimestamp = Long.parseLong(waitStartTimestamp);
        } catch (NumberFormatException nfe) {
            logger.error("{} has an invalid value '{}' on FlowFile {}", new Object[] { WAIT_START_TIMESTAMP, waitStartTimestamp, flowFile });
            transferToFailure.accept(flowFile);
            continue;
        }
        // check for expiration
        long expirationDuration = context.getProperty(EXPIRATION_DURATION).asTimePeriod(TimeUnit.MILLISECONDS);
        long now = System.currentTimeMillis();
        if (now > (lWaitStartTimestamp + expirationDuration)) {
            logger.info("FlowFile {} expired after {}ms", new Object[] { flowFile, (now - lWaitStartTimestamp) });
            getFlowFilesFor.apply(REL_EXPIRED).add(flowFile);
            continue;
        }
        // If there's no signal yet, then we don't have to evaluate target counts. Return immediately.
        if (signal == null) {
            if (logger.isDebugEnabled()) {
                logger.debug("No release signal found for {} on FlowFile {} yet", new Object[] { signalId, flowFile });
            }
            getFlowFilesFor.apply(REL_WAIT).add(flowFile);
            continue;
        }
        // Fix target counter name and count from current FlowFile, if those are not set yet.
        if (candidates.isEmpty()) {
            targetCounterName = context.getProperty(SIGNAL_COUNTER_NAME).evaluateAttributeExpressions(flowFile).getValue();
            try {
                targetCount = Long.valueOf(context.getProperty(TARGET_SIGNAL_COUNT).evaluateAttributeExpressions(flowFile).getValue());
            } catch (final NumberFormatException e) {
                transferToFailure.accept(flowFile);
                logger.error("Failed to parse targetCount when processing {} due to {}", new Object[] { flowFile, e }, e);
                continue;
            }
            try {
                releasableFlowFileCount = Integer.valueOf(context.getProperty(RELEASABLE_FLOWFILE_COUNT).evaluateAttributeExpressions(flowFile).getValue());
            } catch (final NumberFormatException e) {
                transferToFailure.accept(flowFile);
                logger.error("Failed to parse releasableFlowFileCount when processing {} due to {}", new Object[] { flowFile, e }, e);
                continue;
            }
        }
        // FlowFile is now validated and added to candidates.
        candidates.add(flowFile);
    }
    boolean waitCompleted = false;
    boolean waitProgressed = false;
    if (signal != null && !candidates.isEmpty()) {
        if (releasableFlowFileCount > 0) {
            signal.releaseCandidates(targetCounterName, targetCount, releasableFlowFileCount, candidates, released -> getFlowFilesFor.apply(REL_SUCCESS).addAll(released), waiting -> getFlowFilesFor.apply(REL_WAIT).addAll(waiting));
            waitCompleted = signal.getTotalCount() == 0 && signal.getReleasableCount() == 0;
            waitProgressed = !getFlowFilesFor.apply(REL_SUCCESS).isEmpty();
        } else {
            boolean reachedTargetCount = StringUtils.isBlank(targetCounterName) ? signal.isTotalCountReached(targetCount) : signal.isCountReached(targetCounterName, targetCount);
            if (reachedTargetCount) {
                getFlowFilesFor.apply(REL_SUCCESS).addAll(candidates);
            } else {
                getFlowFilesFor.apply(REL_WAIT).addAll(candidates);
            }
        }
    }
    // Transfer FlowFiles.
    processedFlowFiles.entrySet().forEach(transferFlowFiles);
    // Update signal if needed.
    try {
        if (waitCompleted) {
            protocol.complete(signalId);
        } else if (waitProgressed) {
            protocol.replace(signal);
        }
    } catch (final IOException e) {
        session.rollback();
        throw new ProcessException(String.format("Unable to communicate with cache while updating %s due to %s", signalId, e), e);
    }
}
Also used : StandardValidators(org.apache.nifi.processor.util.StandardValidators) FlowFileFilter(org.apache.nifi.processor.FlowFileFilter) CapabilityDescription(org.apache.nifi.annotation.documentation.CapabilityDescription) ResultType(org.apache.nifi.expression.AttributeExpression.ResultType) HashMap(java.util.HashMap) EventDriven(org.apache.nifi.annotation.behavior.EventDriven) ACCEPT_AND_CONTINUE(org.apache.nifi.processor.FlowFileFilter.FlowFileFilterResult.ACCEPT_AND_CONTINUE) ComponentLog(org.apache.nifi.logging.ComponentLog) AtomicReference(java.util.concurrent.atomic.AtomicReference) Function(java.util.function.Function) Supplier(java.util.function.Supplier) StringUtils(org.apache.commons.lang3.StringUtils) PropertyDescriptor(org.apache.nifi.components.PropertyDescriptor) ProcessException(org.apache.nifi.processor.exception.ProcessException) ArrayList(java.util.ArrayList) PropertyValue(org.apache.nifi.components.PropertyValue) HashSet(java.util.HashSet) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) WritesAttributes(org.apache.nifi.annotation.behavior.WritesAttributes) Relationship(org.apache.nifi.processor.Relationship) Map(java.util.Map) Requirement(org.apache.nifi.annotation.behavior.InputRequirement.Requirement) ACCEPT_AND_TERMINATE(org.apache.nifi.processor.FlowFileFilter.FlowFileFilterResult.ACCEPT_AND_TERMINATE) AtomicDistributedMapCacheClient(org.apache.nifi.distributed.cache.client.AtomicDistributedMapCacheClient) Signal(org.apache.nifi.processors.standard.WaitNotifyProtocol.Signal) FlowFile(org.apache.nifi.flowfile.FlowFile) ProcessContext(org.apache.nifi.processor.ProcessContext) Set(java.util.Set) IOException(java.io.IOException) ProcessSession(org.apache.nifi.processor.ProcessSession) WritesAttribute(org.apache.nifi.annotation.behavior.WritesAttribute) SeeAlso(org.apache.nifi.annotation.documentation.SeeAlso) AllowableValue(org.apache.nifi.components.AllowableValue) Collectors(java.util.stream.Collectors) REJECT_AND_CONTINUE(org.apache.nifi.processor.FlowFileFilter.FlowFileFilterResult.REJECT_AND_CONTINUE) TimeUnit(java.util.concurrent.TimeUnit) Consumer(java.util.function.Consumer) List(java.util.List) InputRequirement(org.apache.nifi.annotation.behavior.InputRequirement) SupportsBatching(org.apache.nifi.annotation.behavior.SupportsBatching) Entry(java.util.Map.Entry) AbstractProcessor(org.apache.nifi.processor.AbstractProcessor) Tags(org.apache.nifi.annotation.documentation.Tags) Collections(java.util.Collections) HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) Signal(org.apache.nifi.processors.standard.WaitNotifyProtocol.Signal) Entry(java.util.Map.Entry) AtomicDistributedMapCacheClient(org.apache.nifi.distributed.cache.client.AtomicDistributedMapCacheClient) ArrayList(java.util.ArrayList) List(java.util.List) FlowFile(org.apache.nifi.flowfile.FlowFile) PropertyValue(org.apache.nifi.components.PropertyValue) AtomicReference(java.util.concurrent.atomic.AtomicReference) IOException(java.io.IOException) ComponentLog(org.apache.nifi.logging.ComponentLog) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) ProcessException(org.apache.nifi.processor.exception.ProcessException) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) Relationship(org.apache.nifi.processor.Relationship)

Aggregations

IOException (java.io.IOException)1 ArrayList (java.util.ArrayList)1 Collections (java.util.Collections)1 HashMap (java.util.HashMap)1 HashSet (java.util.HashSet)1 List (java.util.List)1 Map (java.util.Map)1 Entry (java.util.Map.Entry)1 Set (java.util.Set)1 TimeUnit (java.util.concurrent.TimeUnit)1 AtomicInteger (java.util.concurrent.atomic.AtomicInteger)1 AtomicReference (java.util.concurrent.atomic.AtomicReference)1 Consumer (java.util.function.Consumer)1 Function (java.util.function.Function)1 Supplier (java.util.function.Supplier)1 Collectors (java.util.stream.Collectors)1 StringUtils (org.apache.commons.lang3.StringUtils)1 EventDriven (org.apache.nifi.annotation.behavior.EventDriven)1 InputRequirement (org.apache.nifi.annotation.behavior.InputRequirement)1 Requirement (org.apache.nifi.annotation.behavior.InputRequirement.Requirement)1