Repository: apache/flink
File: flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/source/reader/SourceReaderBase.java (Java)

Code before:

        numRecordsInCounter = context.metricGroup().getIOMetricGroup().getNumRecordsInCounter();
        rateLimiter =
                rateLimiterStrategy == null
                        ? null
                        : rateLimiterStrategy.createRateLimiter(context.currentParallelism());
    }

    @Override
    public void start() {}

    @Override
    public InputStatus pollNext(ReaderOutput<T> output) throws Exception {
        // make sure we have a fetch we are working on, or move to the next
        RecordsWithSplitIds<E> recordsWithSplitId = this.currentFetch;
        if (recordsWithSplitId == null) {
            recordsWithSplitId = getNextFetch(output);
            if (recordsWithSplitId == null) {
                return trace(finishedOrAvailableLater());
            }
        }

        // we need to loop here, because we may have to go across splits
        while (true) {
            // Check if the previous record count reached the limit of rateLimiter.
            if (rateLimitPermissionFuture != null && !rateLimitPermissionFuture.isDone()) {
                return trace(InputStatus.MORE_AVAILABLE);
            }
            // Process one record.
            final E record = recordsWithSplitId.nextRecordFromSplit();
            if (record != null) {
                // emit the record.
                numRecordsInCounter.inc(1);
                recordEmitter.emitRecord(record, currentSplitOutput, currentSplitContext.state);
                LOG.trace("Emitted record: {}", record);
                if (rateLimiter != null) {
                    RecordCountingSourceOutputWrapper<T> recordCountingSourceOutputWrapper =
                            (RecordCountingSourceOutputWrapper<T>) currentSplitOutput;
                    if (recordCountingSourceOutputWrapper.getRecordCount() > 0) {
                        // Acquire permit from rateLimiter.
                        rateLimitPermissionFuture =
                                rateLimiter
                                        .acquire(recordCountingSourceOutputWrapper.getRecordCount())
                                        .toCompletableFuture();
                    }
                    recordCountingSourceOutputWrapper.resetRecordCount();
                }
                // We always emit MORE_AVAILABLE here, even though we do not strictly know whether
                // more is available. If nothing more is available, the next invocation will find
                // this out and return the correct status.
                // That means we emit the occasional 'false positive' for availability, but this

Review comment (security):
In current implementation, we need to check here for **each record** even rateLimitPermissionFuture is null , I believe this will introduce a performance regression. We can avoid this regression as we've known that `rateLimiterStrategy` nullability in constructor. 

Code after:
        this.config = config;
        this.context = context;
        this.noMoreSplitsAssignment = false;
        this.eofRecordEvaluator = eofRecordEvaluator;
        this.rateLimitingEnabled = rateLimiterStrategy != null;
        numRecordsInCounter = context.metricGroup().getIOMetricGroup().getNumRecordsInCounter();

        rateLimiter =
                rateLimitingEnabled
                        ? rateLimiterStrategy.createRateLimiter(context.currentParallelism())
                        : null;
        LOG.info(
                "Rate limiting of SourceReader is {}",
                rateLimitingEnabled ? "enabled" : "disabled");
        rateLimitPermissionFuture = CompletableFuture.completedFuture(null);
    }

    @Override
    public void start() {}

    @Override
    public InputStatus pollNext(ReaderOutput<T> output) throws Exception {
        // make sure we have a fetch we are working on, or move to the next
        RecordsWithSplitIds<E> recordsWithSplitId = this.currentFetch;
        if (recordsWithSplitId == null) {
            recordsWithSplitId = getNextFetch(output);
            if (recordsWithSplitId == null) {
                return trace(finishedOrAvailableLater());
            }
        }
        if (rateLimitingEnabled) {
            return pollNextWithRateLimiting(recordsWithSplitId, output);
        } else {
            return pollNextWithoutRateLimiting(recordsWithSplitId, output);
        }
    }

    private InputStatus pollNextWithoutRateLimiting(
            RecordsWithSplitIds<E> recordsWithSplitId, ReaderOutput<T> output) throws Exception {
        // we need to loop here, in case we may have to go across splits
        while (true) {
            // Process one record.
            final E record = recordsWithSplitId.nextRecordFromSplit();
            if (record != null) {
                // emit the record.
                numRecordsInCounter.inc(1);
                recordEmitter.emitRecord(record, currentSplitOutput, currentSplitContext.state);
                LOG.trace("Emitted record: {}", record);
                // We always emit MORE_AVAILABLE here, even though we do not strictly know whether
                // more is available. If nothing more is available, the next invocation will find
                // this out and return the correct status.