/*
 * 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.processors.standard;

import org.apache.nifi.annotation.behavior.InputRequirement;
import org.apache.nifi.annotation.behavior.InputRequirement.Requirement;
import org.apache.nifi.annotation.behavior.SideEffectFree;
import org.apache.nifi.annotation.behavior.TriggerSerially;
import org.apache.nifi.annotation.documentation.CapabilityDescription;
import org.apache.nifi.annotation.documentation.Tags;
import org.apache.nifi.annotation.lifecycle.OnScheduled;
import org.apache.nifi.components.AllowableValue;
import org.apache.nifi.components.PropertyDescriptor;
import org.apache.nifi.components.ValidationContext;
import org.apache.nifi.components.ValidationResult;
import org.apache.nifi.expression.ExpressionLanguageScope;
import org.apache.nifi.flowfile.FlowFile;
import org.apache.nifi.logging.ComponentLog;
import org.apache.nifi.processor.AbstractProcessor;
import org.apache.nifi.processor.DataUnit;
import org.apache.nifi.processor.FlowFileFilter;
import org.apache.nifi.processor.ProcessContext;
import org.apache.nifi.processor.ProcessSession;
import org.apache.nifi.processor.ProcessorInitializationContext;
import org.apache.nifi.processor.Relationship;
import org.apache.nifi.processor.exception.ProcessException;
import org.apache.nifi.processor.util.StandardValidators;
import org.apache.nifi.util.timebuffer.LongEntityAccess;
import org.apache.nifi.util.timebuffer.TimedBuffer;
import org.apache.nifi.util.timebuffer.TimestampedLong;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.LongSupplier;
import java.util.regex.Pattern;

@SideEffectFree
@TriggerSerially
@InputRequirement(Requirement.INPUT_REQUIRED)
@Tags({"rate control", "throttle", "rate", "throughput"})
@CapabilityDescription("Controls the rate at which data is transferred to follow-on processors."
        + " If you configure a very small Time Duration, then the accuracy of the throttle gets worse."
        + " You can improve this accuracy by decreasing the Yield Duration, at the expense of more Tasks given to the processor.")
public class ControlRate extends AbstractProcessor {

    public static final String DATA_RATE = "data rate";
    public static final String FLOWFILE_RATE = "flowfile count";
    public static final String ATTRIBUTE_RATE = "attribute value";
    public static final String DATA_OR_FLOWFILE_RATE = "data rate or flowfile count";

    public static final AllowableValue DATA_RATE_VALUE = new AllowableValue(DATA_RATE, DATA_RATE,
            "Rate is controlled by counting bytes transferred per time duration.");
    public static final AllowableValue FLOWFILE_RATE_VALUE = new AllowableValue(FLOWFILE_RATE, FLOWFILE_RATE,
            "Rate is controlled by counting FlowFiles transferred per time duration");
    public static final AllowableValue ATTRIBUTE_RATE_VALUE = new AllowableValue(ATTRIBUTE_RATE, ATTRIBUTE_RATE,
            "Rate is controlled by accumulating the value of a specified attribute that is transferred per time duration");
    public static final AllowableValue DATA_OR_FLOWFILE_RATE_VALUE = new AllowableValue(DATA_OR_FLOWFILE_RATE, DATA_OR_FLOWFILE_RATE,
            "Rate is controlled by counting bytes and FlowFiles transferred per time duration; if either threshold is met, throttling is enforced");

    // based on testing to balance commits and 10,000 FF swap limit
    public static final int MAX_FLOW_FILES_PER_BATCH = 1000;
    private static final long DEFAULT_ACCRUAL_COUNT = -1L;

    public static final PropertyDescriptor RATE_CONTROL_CRITERIA = new PropertyDescriptor.Builder()
            .name("Rate Control Criteria")
            .displayName("Rate Control Criteria")
            .description("Indicates the criteria that is used to control the throughput rate. Changing this value resets the rate counters.")
            .required(true)
            .allowableValues(DATA_RATE_VALUE, FLOWFILE_RATE_VALUE, ATTRIBUTE_RATE_VALUE, DATA_OR_FLOWFILE_RATE_VALUE)
            .defaultValue(DATA_RATE)
            .build();
    public static final PropertyDescriptor MAX_RATE = new PropertyDescriptor.Builder()
            .name("Maximum Rate")
            .displayName("Maximum Rate")
            .description("The maximum rate at which data should pass through this processor. The format of this property is expected to be a "
                    + "positive integer, or a Data Size (such as '1 MB') if Rate Control Criteria is set to 'data rate'.")
            .required(false)
            .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) // validated in customValidate b/c dependent on Rate Control Criteria
            .dependsOn(RATE_CONTROL_CRITERIA, DATA_RATE_VALUE, FLOWFILE_RATE_VALUE, ATTRIBUTE_RATE_VALUE)
            .build();
    public static final PropertyDescriptor MAX_DATA_RATE = new PropertyDescriptor.Builder()
            .name("Maximum Data Rate")
            .displayName("Maximum Data Rate")
            .description("The maximum rate at which data should pass through this processor. The format of this property is expected to be a "
                    + "Data Size (such as '1 MB') representing bytes per Time Duration.")
            .required(false)
            .addValidator(StandardValidators.DATA_SIZE_VALIDATOR)
            .dependsOn(RATE_CONTROL_CRITERIA, DATA_OR_FLOWFILE_RATE)
            .build();

    public static final PropertyDescriptor MAX_COUNT_RATE = new PropertyDescriptor.Builder()
            .name("Maximum FlowFile Rate")
            .displayName("Maximum FlowFile Rate")
            .description("The maximum rate at which FlowFiles should pass through this processor. The format of this property is expected to be a "
                    + "positive integer representing FlowFiles count per Time Duration")
            .required(false)
            .addValidator(StandardValidators.POSITIVE_LONG_VALIDATOR)
            .dependsOn(RATE_CONTROL_CRITERIA, DATA_OR_FLOWFILE_RATE)
            .build();

    public static final PropertyDescriptor RATE_CONTROL_ATTRIBUTE_NAME = new PropertyDescriptor.Builder()
            .name("Rate Controlled Attribute")
            .description("The name of an attribute whose values build toward the rate limit if Rate Control Criteria is set to 'attribute value'. "
                    + "The value of the attribute referenced by this property must be a positive long, or the FlowFile will be routed to failure. "
                    + "This value is ignored if Rate Control Criteria is not set to 'attribute value'. Changing this value resets the rate counters.")
            .required(false)
            .addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
            .expressionLanguageSupported(ExpressionLanguageScope.NONE)
            .dependsOn(RATE_CONTROL_CRITERIA, ATTRIBUTE_RATE)
            .build();
    public static final PropertyDescriptor TIME_PERIOD = new PropertyDescriptor.Builder()
            .name("Time Duration")
            .description("The amount of time to which the Maximum Rate pertains. Changing this value resets the rate counters.")
            .required(true)
            .addValidator(StandardValidators.createTimePeriodValidator(1, TimeUnit.SECONDS, Integer.MAX_VALUE, TimeUnit.SECONDS))
            .defaultValue("1 min")
            .build();
    public static final PropertyDescriptor GROUPING_ATTRIBUTE_NAME = new PropertyDescriptor.Builder()
            .name("Grouping Attribute")
            .description("By default, a single \"throttle\" is used for all FlowFiles. If this value is specified, a separate throttle is used for "
                    + "each value specified by the attribute with this name. Changing this value resets the rate counters.")
            .required(false)
            .addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
            .expressionLanguageSupported(ExpressionLanguageScope.NONE)
            .build();

    public static final Relationship REL_SUCCESS = new Relationship.Builder()
            .name("success")
            .description("FlowFiles are transferred to this relationship under normal conditions")
            .build();
    public static final Relationship REL_FAILURE = new Relationship.Builder()
            .name("failure")
            .description("FlowFiles will be routed to this relationship if they are missing a necessary Rate Controlled Attribute or the attribute is not in the expected format")
            .build();

    private static final Pattern POSITIVE_LONG_PATTERN = Pattern.compile("0*[1-9][0-9]*");
    private static final String DEFAULT_GROUP_ATTRIBUTE = ControlRate.class.getName() + "###____DEFAULT_GROUP_ATTRIBUTE___###";

    private List<PropertyDescriptor> properties;
    private Set<Relationship> relationships;

    private final ConcurrentMap<String, Throttle> dataThrottleMap = new ConcurrentHashMap<>();
    private final ConcurrentMap<String, Throttle> countThrottleMap = new ConcurrentHashMap<>();
    private final AtomicLong lastThrottleClearTime = new AtomicLong(getCurrentTimeMillis());
    private volatile String rateControlCriteria = null;
    private volatile String rateControlAttribute = null;
    private volatile String maximumRateStr = null;
    private volatile String maximumCountRateStr = null;
    private volatile String groupingAttributeName = null;
    private volatile int timePeriodSeconds = 1;

    @Override
    protected void init(final ProcessorInitializationContext context) {
        final List<PropertyDescriptor> properties = new ArrayList<>();
        properties.add(RATE_CONTROL_CRITERIA);
        properties.add(TIME_PERIOD);
        properties.add(MAX_RATE);
        properties.add(MAX_DATA_RATE);
        properties.add(MAX_COUNT_RATE);
        properties.add(RATE_CONTROL_ATTRIBUTE_NAME);
        properties.add(GROUPING_ATTRIBUTE_NAME);
        this.properties = Collections.unmodifiableList(properties);

        final Set<Relationship> relationships = new HashSet<>();
        relationships.add(REL_SUCCESS);
        relationships.add(REL_FAILURE);
        this.relationships = Collections.unmodifiableSet(relationships);
    }

    @Override
    protected List<PropertyDescriptor> getSupportedPropertyDescriptors() {
        return properties;
    }

    @Override
    public Set<Relationship> getRelationships() {
        return relationships;
    }

    @Override
    protected Collection<ValidationResult> customValidate(final ValidationContext context) {
        final List<ValidationResult> validationResults = new ArrayList<>(super.customValidate(context));

        switch (context.getProperty(RATE_CONTROL_CRITERIA).getValue().toLowerCase()) {
            case DATA_OR_FLOWFILE_RATE:
                // enforce validators to be sure properties are configured; they are only required for DATA_OR_FLOWFILE_RATE criteria
                validationResults.add(StandardValidators.DATA_SIZE_VALIDATOR.validate(MAX_DATA_RATE.getDisplayName(), context.getProperty(MAX_DATA_RATE).getValue(), context));
                validationResults.add(StandardValidators.POSITIVE_LONG_VALIDATOR.validate(MAX_COUNT_RATE.getDisplayName(), context.getProperty(MAX_COUNT_RATE).getValue(), context));
                break;
            case DATA_RATE:
                validationResults.add(StandardValidators.DATA_SIZE_VALIDATOR.validate("Maximum Rate", context.getProperty(MAX_RATE).getValue(), context));
                break;

            case ATTRIBUTE_RATE:
                final String rateAttr = context.getProperty(RATE_CONTROL_ATTRIBUTE_NAME).getValue();
                if (rateAttr == null) {
                    validationResults.add(new ValidationResult.Builder()
                            .subject(RATE_CONTROL_ATTRIBUTE_NAME.getName())
                            .explanation("property must be set if using <Rate Control Criteria> of 'attribute value'")
                            .build());
                }
            case FLOWFILE_RATE:
                validationResults.add(StandardValidators.POSITIVE_LONG_VALIDATOR.validate("Maximum Rate", context.getProperty(MAX_RATE).getValue(), context));
                break;
            default:
                // no custom validation required
                break;
        }

        return validationResults;
    }

    @Override
    public void onPropertyModified(final PropertyDescriptor descriptor, final String oldValue, final String newValue) {
        super.onPropertyModified(descriptor, oldValue, newValue);

        if (descriptor.equals(RATE_CONTROL_CRITERIA)
                || descriptor.equals(RATE_CONTROL_ATTRIBUTE_NAME)
                || descriptor.equals(GROUPING_ATTRIBUTE_NAME)
                || descriptor.equals(TIME_PERIOD)) {
            // if the criteria that is being used to determine limits/throttles is changed, we must clear our throttle map.
            dataThrottleMap.clear();
            countThrottleMap.clear();
        } else if (descriptor.equals(MAX_RATE) || descriptor.equals(MAX_DATA_RATE)) {
            // MAX_RATE could affect either throttle map; MAX_DATA_RATE only affects data throttle map
            final long newRate;
            if (newValue != null) {
                if (DataUnit.DATA_SIZE_PATTERN.matcher(newValue.toUpperCase()).matches()) {
                    newRate = DataUnit.parseDataSize(newValue, DataUnit.B).longValue();
                } else {
                    newRate = Long.parseLong(newValue);
                }
                if (dataThrottleRequired()) {
                    for (final Throttle throttle : dataThrottleMap.values()) {
                        throttle.setMaxRate(newRate);
                    }
                }
                if (countThrottleRequired()) {
                    for (final Throttle throttle : countThrottleMap.values()) {
                        throttle.setMaxRate(newRate);
                    }
                }
            }
        } else if (descriptor.equals(MAX_COUNT_RATE)) {
            // MAX_COUNT_RATE only affects count throttle map
            long newRate;
            try {
                newRate = Long.parseLong(newValue);
            } catch (NumberFormatException nfe) {
                newRate = -1;
            }
            for (final Throttle throttle : countThrottleMap.values()) {
                throttle.setMaxRate(newRate);
            }
        }
    }

    @OnScheduled
    public void onScheduled(final ProcessContext context) {
        rateControlCriteria = context.getProperty(RATE_CONTROL_CRITERIA).getValue().toLowerCase();
        rateControlAttribute = context.getProperty(RATE_CONTROL_ATTRIBUTE_NAME).getValue();
        if (dataThrottleRequired()) {
            // Use MAX_DATA_RATE only for DATA_OR_FLOWFILE_RATE criteria
            maximumRateStr = rateControlCriteria.equals(DATA_OR_FLOWFILE_RATE)
                    ? context.getProperty(MAX_DATA_RATE).getValue().toUpperCase() : context.getProperty(MAX_RATE).getValue().toUpperCase();
        }
        if (countThrottleRequired()) {
            // Use MAX_COUNT_RATE only for DATA_OR_FLOWFILE_RATE criteria
            maximumCountRateStr = rateControlCriteria.equals(DATA_OR_FLOWFILE_RATE)
                    ? context.getProperty(MAX_COUNT_RATE).getValue() : context.getProperty(MAX_RATE).getValue();
        }
        groupingAttributeName = context.getProperty(GROUPING_ATTRIBUTE_NAME).getValue();
        timePeriodSeconds = context.getProperty(TIME_PERIOD).asTimePeriod(TimeUnit.SECONDS).intValue();
    }

    @Override
    public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
        List<FlowFile> flowFiles = session.get(new ThrottleFilter(MAX_FLOW_FILES_PER_BATCH, this::getCurrentTimeMillis));
        if (flowFiles.isEmpty()) {
            context.yield();
            return;
        }

        // Periodically clear any Throttle that has not been used in more than 2 throttling periods
        final long lastClearTime = lastThrottleClearTime.get();
        final long throttleExpirationMillis = getCurrentTimeMillis() - 2 * context.getProperty(TIME_PERIOD).asTimePeriod(TimeUnit.MILLISECONDS);
        if (lastClearTime < throttleExpirationMillis) {
            if (lastThrottleClearTime.compareAndSet(lastClearTime, getCurrentTimeMillis())) {
                final Set<Map.Entry<String, Throttle>> throttleSet = new HashSet<>();
                if (dataThrottleRequired()) {
                    throttleSet.addAll(dataThrottleMap.entrySet());
                }
                if (countThrottleRequired()) {
                    throttleSet.addAll(countThrottleMap.entrySet());
                }
                final Iterator<Map.Entry<String, Throttle>> itr = throttleSet.iterator();
                while (itr.hasNext()) {
                    final Map.Entry<String, Throttle> entry = itr.next();
                    final Throttle throttle = entry.getValue();
                    if (throttle.tryLock()) {
                        try {
                            if (throttle.lastUpdateTime() < lastClearTime) {
                                itr.remove();
                            }
                        } finally {
                            throttle.unlock();
                        }
                    }
                }
            }
        }

        final ComponentLog logger = getLogger();
        for (FlowFile flowFile : flowFiles) {
            // call this to capture potential error
            if (isAccrualPossible(flowFile)) {
                logger.info("transferring {} to 'success'", flowFile);
                session.transfer(flowFile, REL_SUCCESS);
            } else {
                logger.error("Routing {} to 'failure' due to missing or invalid attribute", flowFile);
                session.transfer(flowFile, REL_FAILURE);
            }
        }
    }

    /**
     * Get current time in milliseconds
     *
     * @return Current time in milliseconds from System
     */
    protected long getCurrentTimeMillis() {
        return System.currentTimeMillis();
    }

    /*
     * Determine if the accrual amount is valid for the type of throttle being applied. For example, if throttling based on
     * flowfile attribute, the specified attribute must be present and must be a long integer.
     */
    private boolean isAccrualPossible(FlowFile flowFile) {
        if (rateControlCriteria.equals(ATTRIBUTE_RATE)) {
            final String attributeValue = flowFile.getAttribute(rateControlAttribute);
            return attributeValue != null && POSITIVE_LONG_PATTERN.matcher(attributeValue).matches();
        }
        return true;
    }

    /*
     * Determine the amount this FlowFile will incur against the maximum allowed rate.
     * This is applicable to data size accrual only
     */
    private long getDataSizeAccrual(FlowFile flowFile) {
        return flowFile.getSize();
    }

    /*
     * Determine the amount this FlowFile will incur against the maximum allowed rate.
     * This is applicable to counting accruals, flowfiles or attributes
     */
    private long getCountAccrual(FlowFile flowFile) {
        long rateValue = DEFAULT_ACCRUAL_COUNT;
        if (rateControlCriteria.equals(FLOWFILE_RATE) || rateControlCriteria.equals(DATA_OR_FLOWFILE_RATE)) {
            rateValue = 1;
        }
        if (rateControlCriteria.equals(ATTRIBUTE_RATE)) {
            final String attributeValue = flowFile.getAttribute(rateControlAttribute);
            if (attributeValue == null) {
                return DEFAULT_ACCRUAL_COUNT;
            }

            if (!POSITIVE_LONG_PATTERN.matcher(attributeValue).matches()) {
                return DEFAULT_ACCRUAL_COUNT;
            }
            rateValue = Long.parseLong(attributeValue);
        }
        return rateValue;
    }

    private boolean dataThrottleRequired() {
        return rateControlCriteria != null && (rateControlCriteria.equals(DATA_RATE) || rateControlCriteria.equals(DATA_OR_FLOWFILE_RATE));
    }

    private boolean countThrottleRequired() {
        return rateControlCriteria != null && (rateControlCriteria.equals(FLOWFILE_RATE) || rateControlCriteria.equals(ATTRIBUTE_RATE) || rateControlCriteria.equals(DATA_OR_FLOWFILE_RATE));
    }

    private static class Throttle extends ReentrantLock {

        private final AtomicLong maxRate = new AtomicLong(1L);
        private final long timePeriodMillis;
        private final TimedBuffer<TimestampedLong> timedBuffer;
        private final ComponentLog logger;
        private final LongSupplier currentTimeSupplier;

        private volatile long penalizationPeriod = 0;
        private volatile long penalizationExpired = 0;
        private volatile long lastUpdateTime;

        private Throttle(final int timePeriod, final TimeUnit unit, final ComponentLog logger, final LongSupplier currentTimeSupplier) {
            this.timePeriodMillis = TimeUnit.MILLISECONDS.convert(timePeriod, unit);
            this.timedBuffer = new TimedBuffer<>(unit, timePeriod, new LongEntityAccess(), currentTimeSupplier);
            this.logger = logger;
            this.currentTimeSupplier = currentTimeSupplier;
        }

        public void setMaxRate(final long maxRate) {
            this.maxRate.set(maxRate);
        }

        public long lastUpdateTime() {
            return lastUpdateTime;
        }

        public boolean tryAdd(final long value) {
            // value should never be negative, but if it is return immediately
            if (value < 0) {
                return false;
            }
            final long now = currentTimeSupplier.getAsLong();
            if (penalizationExpired > now) {
                return false;
            }

            final long maxRateValue = maxRate.get();

            final TimestampedLong sum = timedBuffer.getAggregateValue(timePeriodMillis);
            if (sum != null && sum.getValue() >= maxRateValue) {
                if (logger.isDebugEnabled()) {
                    logger.debug("current sum for throttle is {} at time {}, so not allowing rate of {} through", sum.getValue(), sum.getTimestamp(), value);
                }
                return false;
            }

            // Implement the Throttle penalization based on how much extra 'amountOver' was allowed through
            if (penalizationPeriod > 0) {
                if (logger.isDebugEnabled()) {
                    logger.debug("Starting Throttle penalization, expiring {} milliseconds from now", penalizationPeriod);
                }
                penalizationExpired = now + penalizationPeriod;
                penalizationPeriod = 0;
                return false;
            }

            if (logger.isDebugEnabled()) {
                logger.debug("current sum for throttle is {} at time {}, so allowing rate of {} through",
                        sum == null ? 0 : sum.getValue(), sum == null ? 0 : sum.getTimestamp(), value);
            }

            final long transferred = timedBuffer.add(new TimestampedLong(value)).getValue();
            if (transferred > maxRateValue) {
                final long amountOver = transferred - maxRateValue;
                // determine how long it should take to transfer 'amountOver' and 'penalize' the Throttle for that long
                final double pct = (double) amountOver / (double) maxRateValue;
                this.penalizationPeriod = (long) (timePeriodMillis * pct);

                if (logger.isDebugEnabled()) {
                    logger.debug("allowing rate of {} through but penalizing Throttle for {} milliseconds", value, penalizationPeriod);
                }
            }

            lastUpdateTime = now;
            return true;
        }
    }

    private class ThrottleFilter implements FlowFileFilter {

        private final int flowFilesPerBatch;
        private final LongSupplier currentTimeSupplier;
        private int flowFilesInBatch = 0;

        ThrottleFilter(final int maxFFPerBatch, final LongSupplier currentTimeSupplier) {
            this.flowFilesPerBatch = maxFFPerBatch;
            this.currentTimeSupplier = currentTimeSupplier;
        }

        @Override
        public FlowFileFilterResult filter(FlowFile flowFile) {
            if (!isAccrualPossible(flowFile)) {
                // this FlowFile is invalid for this configuration so let the processor deal with it
                return FlowFileFilterResult.ACCEPT_AND_TERMINATE;
            }

            String groupName = (groupingAttributeName == null) ? DEFAULT_GROUP_ATTRIBUTE : flowFile.getAttribute(groupingAttributeName);

            // the flow file may not have the required attribute: in this case it is considered part
            // of the DEFAULT_GROUP_ATTRIBUTE
            if (groupName == null) {
                groupName = DEFAULT_GROUP_ATTRIBUTE;
            }

            Throttle dataThrottle = dataThrottleMap.get(groupName);
            Throttle countThrottle = countThrottleMap.get(groupName);

            boolean dataThrottlingActive = false;
            if (dataThrottleRequired()) {
                if (dataThrottle == null) {
                    dataThrottle = new Throttle(timePeriodSeconds, TimeUnit.SECONDS, getLogger(), currentTimeSupplier);
                    dataThrottle.setMaxRate(DataUnit.parseDataSize(maximumRateStr, DataUnit.B).longValue());
                    dataThrottleMap.put(groupName, dataThrottle);
                }

                dataThrottle.lock();
                try {
                    if (dataThrottle.tryAdd(getDataSizeAccrual(flowFile))) {
                        flowFilesInBatch++;
                        if (flowFilesInBatch>= flowFilesPerBatch) {
                            flowFilesInBatch = 0;
                            return FlowFileFilterResult.ACCEPT_AND_TERMINATE;
                        } else {
                            // only accept flowfile if additional count throttle does not need to run
                            if (!countThrottleRequired()) {
                                return FlowFileFilterResult.ACCEPT_AND_CONTINUE;
                            }
                        }
                    } else {
                        dataThrottlingActive = true;
                    }
                } finally {
                    dataThrottle.unlock();
                }
            }

            // continue processing count throttle only if required and if data throttle is not already limiting flowfiles
            if (countThrottleRequired() && !dataThrottlingActive) {
                if (countThrottle == null) {
                    countThrottle = new Throttle(timePeriodSeconds, TimeUnit.SECONDS, getLogger(), currentTimeSupplier);
                    countThrottle.setMaxRate(Long.parseLong(maximumCountRateStr));
                    countThrottleMap.put(groupName, countThrottle);
                }
                countThrottle.lock();
                try {
                    if (countThrottle.tryAdd(getCountAccrual(flowFile))) {
                        flowFilesInBatch++;
                        if (flowFilesInBatch>= flowFilesPerBatch) {
                            flowFilesInBatch = 0;
                            return FlowFileFilterResult.ACCEPT_AND_TERMINATE;
                        } else {
                            return FlowFileFilterResult.ACCEPT_AND_CONTINUE;
                        }
                    }
                } finally {
                    countThrottle.unlock();
                }
            }

            // If we are not using a grouping attribute, then no FlowFile will be able to continue on. So we can
            // just TERMINATE the iteration over FlowFiles.
            // However, if we are using a grouping attribute, then another FlowFile in the queue may be able to proceed,
            // so we want to continue our iteration.
            if (groupingAttributeName == null) {
                return FlowFileFilterResult.REJECT_AND_TERMINATE;
            }

            return FlowFileFilterResult.REJECT_AND_CONTINUE;
        }
    }
}
