diff --git a/psc-flink/src/main/java/com/pinterest/flink/connector/psc/source/PscSourceOptions.java b/psc-flink/src/main/java/com/pinterest/flink/connector/psc/source/PscSourceOptions.java index 0f2aabf..e7cf092 100644 --- a/psc-flink/src/main/java/com/pinterest/flink/connector/psc/source/PscSourceOptions.java +++ b/psc-flink/src/main/java/com/pinterest/flink/connector/psc/source/PscSourceOptions.java @@ -58,6 +58,22 @@ public class PscSourceOptions { .defaultValue(true) .withDescription("Whether to commit consuming offset on checkpoint."); + /** + * Total records-per-second budget across all source subtasks. Applied in {@code + * PscTopicUriPartitionSplitReader.fetch} before {@code consumer.poll()}, so MemQ/Kafka + * downloads are paced by the limiter (unlike a downstream map operator). + */ + public static final ConfigOption SCAN_RATE_LIMIT_RECORDS_PER_SECOND = + ConfigOptions.key("scan.rate-limit.records-per-second") + .doubleType() + .noDefaultValue() + .withDescription( + "Optional rate limit for the source in records per second. " + + "When set, each SplitReader acquires permits for the next poll " + + "batch before calling consumer.poll(), dividing the total rate " + + "evenly across source parallelism. If unset, no fetch-side rate " + + "limiting is applied."); + @SuppressWarnings("unchecked") public static T getOption( Properties props, ConfigOption configOption, Function parser) { diff --git a/psc-flink/src/main/java/com/pinterest/flink/connector/psc/source/reader/PscTopicUriPartitionSplitReader.java b/psc-flink/src/main/java/com/pinterest/flink/connector/psc/source/reader/PscTopicUriPartitionSplitReader.java index 9428f28..623627b 100644 --- a/psc-flink/src/main/java/com/pinterest/flink/connector/psc/source/reader/PscTopicUriPartitionSplitReader.java +++ b/psc-flink/src/main/java/com/pinterest/flink/connector/psc/source/reader/PscTopicUriPartitionSplitReader.java @@ -28,7 +28,7 @@ import com.pinterest.psc.consumer.OffsetCommitCallback; import com.pinterest.psc.consumer.PscConsumer; import com.pinterest.psc.consumer.PscConsumerMessage; -import com.pinterest.psc.consumer.PscConsumerMessagesIterable; +import com.pinterest.psc.consumer.PscConsumerPollMessageIterator; import com.pinterest.psc.exception.ClientException; import com.pinterest.psc.exception.consumer.ConsumerException; import com.pinterest.psc.exception.consumer.WakeupException; @@ -40,6 +40,7 @@ import org.apache.flink.connector.base.source.reader.splitreader.SplitReader; import org.apache.flink.connector.base.source.reader.splitreader.SplitsAddition; import org.apache.flink.connector.base.source.reader.splitreader.SplitsChange; +import org.apache.flink.shaded.guava31.com.google.common.util.concurrent.RateLimiter; import org.apache.flink.util.FlinkRuntimeException; import org.apache.flink.util.Preconditions; import org.slf4j.Logger; @@ -52,12 +53,12 @@ import java.util.Collection; import java.util.HashMap; import java.util.HashSet; -import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Properties; import java.util.Set; import java.util.StringJoiner; +import java.util.function.IntConsumer; import java.util.function.Supplier; import java.util.stream.Collectors; @@ -67,6 +68,8 @@ public class PscTopicUriPartitionSplitReader implements SplitReader, PscTopicUriPartitionSplit> { private static final Logger LOG = LoggerFactory.getLogger(PscTopicUriPartitionSplitReader.class); private static final long POLL_TIMEOUT = 10000L; + private static final double MIN_SUBTASK_RATE_LIMIT_QPS = 0.1; + private static final int DEFAULT_POLL_MESSAGES_MAX = 500; private final PscConsumer consumer; private final Map stoppingOffsets; @@ -79,6 +82,18 @@ public class PscTopicUriPartitionSplitReader private final Set emptySplits = new HashSet<>(); private final Properties props; + /** Optional fetch-side rate limiter; null when scan.rate-limit is unset. */ + @Nullable private final RateLimiter fetchRateLimiter; + + /** + * Permits to acquire before the next {@link #fetch()}. Starts at poll.messages.max so the first + * parallel startup fetches are staggered; then tracks the previous poll's emitted count. + */ + private int nextFetchRatePermits; + + /** Partitions finished while streaming the previous poll; unassigned at the start of fetch(). */ + private final List pendingUnassignPartitions = new ArrayList<>(); + public PscTopicUriPartitionSplitReader( Properties props, SourceReaderContext context, @@ -101,24 +116,106 @@ public PscTopicUriPartitionSplitReader( this.consumer = new PscConsumer<>(PscConfigurationUtils.propertiesToPscConfiguration(consumerProps)); this.stoppingOffsets = new HashMap<>(); this.groupId = consumerProps.getProperty(PscConfiguration.PSC_CONSUMER_GROUP_ID); + + int pollMessagesMax = + parsePositiveInt( + props.getProperty(PscConfiguration.PSC_CONSUMER_POLL_MESSAGES_MAX), + DEFAULT_POLL_MESSAGES_MAX); + this.nextFetchRatePermits = pollMessagesMax; + this.fetchRateLimiter = createFetchRateLimiter(props, context.currentParallelism(), pollMessagesMax); + } + + @Nullable + private static RateLimiter createFetchRateLimiter( + Properties props, int parallelism, int pollMessagesMax) { + String rateLimitStr = + props.getProperty(PscSourceOptions.SCAN_RATE_LIMIT_RECORDS_PER_SECOND.key()); + if (rateLimitStr == null || rateLimitStr.isEmpty()) { + return null; + } + double totalRate = Double.parseDouble(rateLimitStr); + if (totalRate <= 0) { + return null; + } + int parallel = Math.max(1, parallelism); + double subtaskRate = totalRate / parallel; + Preconditions.checkArgument( + subtaskRate > MIN_SUBTASK_RATE_LIMIT_QPS, + "Subtask rate limit should be greater than %s QPS. " + + "Current rate: %s records/second divided by %s subtasks = %s records/second per subtask. " + + "Consider increasing the rate limit or decreasing parallelism.", + MIN_SUBTASK_RATE_LIMIT_QPS, + totalRate, + parallel, + subtaskRate); + LOG.info( + "Fetch-side rate limit enabled: {} records/second total, {}/s per subtask " + + "(parallelism={}, initial batch permits={})", + totalRate, + subtaskRate, + parallel, + pollMessagesMax); + return RateLimiter.create(subtaskRate); + } + + private static int parsePositiveInt(@Nullable String value, int defaultValue) { + if (value == null || value.isEmpty()) { + return defaultValue; + } + try { + int parsed = Integer.parseInt(value); + return parsed > 0 ? parsed : defaultValue; + } catch (NumberFormatException e) { + return defaultValue; + } + } + + private void acquireFetchRatePermitsBeforePoll() { + if (fetchRateLimiter == null) { + return; + } + int permits = Math.max(1, nextFetchRatePermits); + fetchRateLimiter.acquire(permits); + } + + private void unassignPendingFinishedPartitions() throws ConsumerException, ConfigurationException { + if (pendingUnassignPartitions.isEmpty()) { + return; + } + pendingUnassignPartitions.forEach(pscSourceReaderMetrics::removeRecordsLagMetric); + unassignPartitions(pendingUnassignPartitions); + pendingUnassignPartitions.clear(); } @Override public RecordsWithSplitIds> fetch() throws IOException { - PscConsumerMessagesIterable consumerMessagesIterable; try { - consumerMessagesIterable = new PscConsumerMessagesIterable<>(consumer.poll(Duration.ofMillis(POLL_TIMEOUT))); + unassignPendingFinishedPartitions(); + } catch (ConsumerException | ConfigurationException e) { + throw new RuntimeException("Failed to unassign finished partitions", e); + } + + // Pace MemQ/Kafka downloads: acquire before poll so fetchObjectToInputStream cannot run + // ahead of the configured record budget. + acquireFetchRatePermitsBeforePoll(); + + PscConsumerPollMessageIterator pollIterator; + try { + pollIterator = consumer.poll(Duration.ofMillis(POLL_TIMEOUT)); } catch (ConsumerException e) { // IllegalStateException will be thrown if the consumer is not assigned any partitions. // This happens if all assigned partitions are invalid or empty (starting offset >= // stopping offset). We just mark empty partitions as finished and return an empty // record container, and this consumer will be closed by SplitFetcherManager. - if (e.getCause() != null && - (e.getCause().getClass().equals(IllegalStateException.class) || e.getCause().getClass().equals(WakeupException.class))) { - LOG.warn("Caught IllegalStateException or WakeupException in poll(), marking partitions as finished", e); + if (e.getCause() != null + && (e.getCause().getClass().equals(IllegalStateException.class) + || e.getCause().getClass().equals(WakeupException.class))) { + LOG.warn( + "Caught IllegalStateException or WakeupException in poll(), marking partitions as finished", + e); + nextFetchRatePermits = 1; PscPartitionSplitRecords recordsBySplits = - new PscPartitionSplitRecords( - PscConsumerMessagesIterable.emptyIterable(), pscSourceReaderMetrics); + PscPartitionSplitRecords.empty(pscSourceReaderMetrics); markEmptySplitsAsFinished(recordsBySplits); return recordsBySplits; } else { @@ -129,48 +226,20 @@ public RecordsWithSplitIds> fetch() throws IO LOG.error("Unrecoverable Exception caught in poll()", e); throw new RuntimeException(e); } + + // Stream records from the poll iterator — do not call asList() / PscConsumerMessagesIterable, + // which materializes every raw payload onto the heap before Flink can emit or backpressure. PscPartitionSplitRecords recordsBySplits = - new PscPartitionSplitRecords(consumerMessagesIterable, pscSourceReaderMetrics); - List finishedPartitions = new ArrayList<>(); - for (TopicUriPartition tp : consumerMessagesIterable.getTopicUriPartitions()) { - long stoppingOffset = getStoppingOffset(tp); - final List> recordsFromPartition = - consumerMessagesIterable.getMessagesForTopicUriPartition(tp); - - if (recordsFromPartition.size() > 0) { - final PscConsumerMessage lastRecord = - recordsFromPartition.get(recordsFromPartition.size() - 1); - - // After processing a record with offset of "stoppingOffset - 1", the split reader - // should not continue fetching because the record with stoppingOffset may not - // exist. Keep polling will just block forever. - if (lastRecord.getMessageId().getOffset() >= stoppingOffset - 1) { - recordsBySplits.setPartitionStoppingOffset(tp, stoppingOffset); - finishSplitAtRecord( - tp, - stoppingOffset, - lastRecord.getMessageId().getOffset(), - finishedPartitions, - recordsBySplits); - } - } - // Track this partition's record lag if it never appears before - pscSourceReaderMetrics.maybeAddRecordsLagMetric(consumer, tp); - } + new PscPartitionSplitRecords( + pollIterator, + stoppingOffsets, + pscSourceReaderMetrics, + pendingUnassignPartitions, + emittedCount -> nextFetchRatePermits = Math.max(1, emittedCount)); markEmptySplitsAsFinished(recordsBySplits); - // Unassign the partitions that has finished. - if (!finishedPartitions.isEmpty()) { - finishedPartitions.forEach(pscSourceReaderMetrics::removeRecordsLagMetric); - try { - unassignPartitions(finishedPartitions); - } catch (ConsumerException | ConfigurationException e) { - throw new RuntimeException("Failed to unassign partitions", e); - } - } - - // Update numBytesIn + // Update numBytesIn (best-effort; streaming path updates as records are read) pscSourceReaderMetrics.updateNumBytesInCounter(); return recordsBySplits; @@ -198,9 +267,9 @@ public void handleSplitsChanges(SplitsChange splitsCh // Assignment. List newPartitionAssignments = new ArrayList<>(); // Starting offsets. - Map partitionsStartingFromSpecifiedOffsets = new HashMap<>(); List partitionsStartingFromEarliest = new ArrayList<>(); List partitionsStartingFromLatest = new ArrayList<>(); + Map partitionsStartingFromSpecifiedOffsets = new HashMap<>(); // Stopping offsets. List partitionsStoppingAtLatest = new ArrayList<>(); Set partitionsStoppingAtCommitted = new HashSet<>(); @@ -301,6 +370,17 @@ PscConsumer consumer() { return consumer; } + @VisibleForTesting + @Nullable + RateLimiter fetchRateLimiter() { + return fetchRateLimiter; + } + + @VisibleForTesting + int nextFetchRatePermits() { + return nextFetchRatePermits; + } + // --------------- private helper method ---------------------- /** @@ -491,21 +571,6 @@ private String createConsumerClientId(Properties props) { return prefix + "-" + subtaskId; } - private void finishSplitAtRecord( - TopicUriPartition tp, - long stoppingOffset, - long currentOffset, - List finishedPartitions, - PscPartitionSplitRecords recordsBySplits) { - LOG.debug( - "{} has reached stopping offset {}, current offset is {}", - tp, - stoppingOffset, - currentOffset); - finishedPartitions.add(tp); - recordsBySplits.addFinishedSplit(PscTopicUriPartitionSplit.toSplitId(tp)); - } - private long getStoppingOffset(TopicUriPartition tp) { return stoppingOffsets.getOrDefault(tp, Long.MAX_VALUE); } @@ -556,52 +621,120 @@ private V retryOnWakeup(Supplier consumerCall, String description) { // ---------------- private helper class ------------------------ + /** + * Streams {@link PscConsumerPollMessageIterator} records into Flink without materializing the + * full poll into a {@code List} via {@code asList()}. + * + *

Records are grouped into splits on the fly: {@link #nextSplit()} starts a partition from + * the peeked message; {@link #nextRecordFromSplit()} emits consecutive messages for that + * partition until the partition changes or the iterator is exhausted. + */ private static class PscPartitionSplitRecords implements RecordsWithSplitIds> { private final Set finishedSplits = new HashSet<>(); - private final Map stoppingOffsets = new HashMap<>(); - private final PscConsumerMessagesIterable consumerMessagesIterable; + private final Map stoppingOffsets; private final PscSourceReaderMetrics metrics; - private final Iterator splitIterator; - private Iterator> recordIterator; + private final List finishedPartitionsForUnassign; + private final IntConsumer onFinishedEmitting; + + @Nullable private final PscConsumerPollMessageIterator pollIterator; + @Nullable private PscConsumerMessage peekedMessage; + private boolean iteratorExhausted; + private TopicUriPartition currentTopicPartition; private Long currentSplitStoppingOffset; private PscSourceReaderMetrics.Offset currentOffsetTracker; + private int emittedCount; + + private boolean reportedEmitCount; private PscPartitionSplitRecords( - PscConsumerMessagesIterable consumerMessagesIterable, PscSourceReaderMetrics metrics) { - this.consumerMessagesIterable = consumerMessagesIterable; - this.splitIterator = consumerMessagesIterable.getTopicUriPartitions().iterator(); + @Nullable PscConsumerPollMessageIterator pollIterator, + Map stoppingOffsets, + PscSourceReaderMetrics metrics, + List finishedPartitionsForUnassign, + IntConsumer onFinishedEmitting) { + this.pollIterator = pollIterator; + this.stoppingOffsets = stoppingOffsets; this.metrics = metrics; + this.finishedPartitionsForUnassign = finishedPartitionsForUnassign; + this.onFinishedEmitting = onFinishedEmitting; + this.iteratorExhausted = pollIterator == null; } - private void setPartitionStoppingOffset( - TopicUriPartition topicUriPartition, long stoppingOffset) { - stoppingOffsets.put(topicUriPartition, stoppingOffset); + private static PscPartitionSplitRecords empty(PscSourceReaderMetrics metrics) { + return new PscPartitionSplitRecords( + null, new HashMap<>(), metrics, new ArrayList<>(), count -> {}); } - private void addFinishedSplit(String splitId) { - finishedSplits.add(splitId); + private void reportEmittedCountOnce() { + if (!reportedEmitCount) { + reportedEmitCount = true; + onFinishedEmitting.accept(emittedCount); + } + } + + private void ensurePeek() { + if (peekedMessage != null || iteratorExhausted) { + return; + } + while (pollIterator != null && pollIterator.hasNext()) { + PscConsumerMessage next = pollIterator.next(); + TopicUriPartition tp = next.getMessageId().getTopicUriPartition(); + // Drop messages for splits already finished in this poll (stopping offset reached). + if (finishedSplits.contains(PscTopicUriPartitionSplit.toSplitId(tp))) { + continue; + } + peekedMessage = next; + return; + } + iteratorExhausted = true; + closeIteratorQuietly(); + reportEmittedCountOnce(); + } + + private void closeIteratorQuietly() { + if (pollIterator == null) { + return; + } + try { + pollIterator.close(); + } catch (IOException e) { + LOG.warn("Failed to close poll message iterator", e); + } + } + + private void maybeFinishSplitAtOffset(long offset) { + if (offset < currentSplitStoppingOffset - 1) { + return; + } + String splitId = PscTopicUriPartitionSplit.toSplitId(currentTopicPartition); + if (finishedSplits.add(splitId)) { + finishedPartitionsForUnassign.add(currentTopicPartition); + LOG.debug( + "{} has reached stopping offset {}, current offset is {}", + currentTopicPartition, + currentSplitStoppingOffset, + offset); + } } @Nullable @Override public String nextSplit() { - if (splitIterator.hasNext()) { - currentTopicPartition = splitIterator.next(); - recordIterator = consumerMessagesIterable.getMessagesForTopicUriPartition(currentTopicPartition).iterator(); - currentSplitStoppingOffset = - stoppingOffsets.getOrDefault(currentTopicPartition, Long.MAX_VALUE); - currentOffsetTracker = metrics.getOffsetTracker(currentTopicPartition); - return currentTopicPartition.toString(); - } else { + ensurePeek(); + if (peekedMessage == null) { currentTopicPartition = null; - recordIterator = null; currentSplitStoppingOffset = null; currentOffsetTracker = null; return null; } + currentTopicPartition = peekedMessage.getMessageId().getTopicUriPartition(); + currentSplitStoppingOffset = + stoppingOffsets.getOrDefault(currentTopicPartition, Long.MAX_VALUE); + currentOffsetTracker = metrics.getOffsetTracker(currentTopicPartition); + return currentTopicPartition.toString(); } @Nullable @@ -611,20 +744,45 @@ public PscConsumerMessage nextRecordFromSplit() { currentTopicPartition, "Make sure nextSplit() did not return null before " + "iterate over the records split."); - if (recordIterator.hasNext()) { - final PscConsumerMessage message = recordIterator.next(); - // Only emit records before stopping offset - if (message.getMessageId().getOffset() < currentSplitStoppingOffset) { - currentOffsetTracker.currentOffset = message.getMessageId().getOffset(); - return message; - } + ensurePeek(); + if (peekedMessage == null) { + return null; } - return null; + TopicUriPartition messageTp = peekedMessage.getMessageId().getTopicUriPartition(); + if (!messageTp.equals(currentTopicPartition)) { + return null; + } + + final PscConsumerMessage message = peekedMessage; + peekedMessage = null; + final long offset = message.getMessageId().getOffset(); + + // Only emit records before the stopping offset (same contract as before). + if (offset >= currentSplitStoppingOffset) { + maybeFinishSplitAtOffset(offset); + return null; + } + + currentOffsetTracker.currentOffset = offset; + emittedCount++; + maybeFinishSplitAtOffset(offset); + return message; } @Override public Set finishedSplits() { return finishedSplits; } + + @Override + public void recycle() { + // Prefer closing any remaining iterator state once Flink is done with this batch. + if (!iteratorExhausted) { + iteratorExhausted = true; + peekedMessage = null; + closeIteratorQuietly(); + } + reportEmittedCountOnce(); + } } } diff --git a/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscConnectorOptions.java b/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscConnectorOptions.java index 5c686a2..8eabb90 100644 --- a/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscConnectorOptions.java +++ b/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscConnectorOptions.java @@ -127,7 +127,9 @@ public class PscConnectorOptions { .noDefaultValue() .withDescription( "Optional rate limit for the source in records per second. " + - "When specified, the source will throttle consumption to not exceed this rate. " + + "When specified, the PSC SplitReader throttles before consumer.poll() " + + "(fetch-side), so backend downloads such as MemQ object fetches are paced " + + "by the limiter — not only record emission after the source. " + "The rate is distributed evenly across all parallel source subtasks. " + "For example, with a rate limit of 1000 and parallelism of 4, each subtask will " + "process approximately 250 records/second. " + diff --git a/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicSource.java b/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicSource.java index 7545c3c..412fb8d 100644 --- a/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicSource.java +++ b/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicSource.java @@ -20,6 +20,7 @@ import com.pinterest.flink.connector.psc.source.PscSource; import com.pinterest.flink.connector.psc.source.PscSourceBuilder; +import com.pinterest.flink.connector.psc.source.PscSourceOptions; import com.pinterest.flink.connector.psc.source.enumerator.initializer.NoStoppingOffsetsInitializer; import com.pinterest.flink.connector.psc.source.enumerator.initializer.OffsetsInitializer; import com.pinterest.flink.connector.psc.source.reader.deserializer.PscRecordDeserializationSchema; @@ -462,14 +463,9 @@ public DataStream produceDataStream( + "parallelism = {}", execEnv.getParallelism()); } - if (isRateLimitingEnabled(rateLimitRecordsPerSecond)) { - String rateLimiterOperatorName = "PscRateLimit-" + tableIdentifier; - resultStream = resultStream - .map(new PscRateLimitMap<>(rateLimitRecordsPerSecond)) - .setParallelism(sourceStream.getParallelism()) - .name(rateLimiterOperatorName) - .uid(rateLimiterOperatorName); - } + // Rate limiting is applied fetch-side inside PscTopicUriPartitionSplitReader + // (before consumer.poll). Do not add a downstream PscRateLimitMap — that would + // throttle only after MemQ/Kafka downloads already hit the heap. if (enableRescale) { resultStream = resultStream.rescale(); } @@ -994,8 +990,21 @@ protected PscSource createPscSource( break; } + Properties sourceProperties = properties; + if (isRateLimitingEnabled(rateLimitRecordsPerSecond)) { + // Copy so we don't mutate shared table properties; SplitReader reads this key. + sourceProperties = new Properties(); + sourceProperties.putAll(properties); + sourceProperties.setProperty( + PscSourceOptions.SCAN_RATE_LIMIT_RECORDS_PER_SECOND.key(), + Double.toString(rateLimitRecordsPerSecond)); + LOG.info( + "Configured fetch-side rate limit: {} records/second (total across all subtasks)", + rateLimitRecordsPerSecond); + } + pscSourceBuilder - .setProperties(properties) + .setProperties(sourceProperties) .setDeserializer(PscRecordDeserializationSchema.of(pscDeserializer)); return pscSourceBuilder.build(); diff --git a/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscRateLimitMap.java b/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscRateLimitMap.java index 5fabe3e..13678be 100644 --- a/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscRateLimitMap.java +++ b/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscRateLimitMap.java @@ -35,6 +35,10 @@ *

When the rate limit is exceeded, the function blocks until permits become available, * emitting metrics to track throttling behavior. * + *

Note: {@link com.pinterest.flink.streaming.connectors.psc.table.PscDynamicSource} + * now applies rate limiting fetch-side inside {@code PscTopicUriPartitionSplitReader} (before + * {@code consumer.poll()}) instead of inserting this map into the operator graph. This class + * remains for unit tests and any callers that still want post-source emission throttling. * * * @param The type of records flowing through this map function diff --git a/psc-flink/src/test/java/com/pinterest/flink/connector/psc/source/reader/PscTopicUriPartitionSplitReaderTest.java b/psc-flink/src/test/java/com/pinterest/flink/connector/psc/source/reader/PscTopicUriPartitionSplitReaderTest.java index 4b4800d..7fa0f74 100644 --- a/psc-flink/src/test/java/com/pinterest/flink/connector/psc/source/reader/PscTopicUriPartitionSplitReaderTest.java +++ b/psc-flink/src/test/java/com/pinterest/flink/connector/psc/source/reader/PscTopicUriPartitionSplitReaderTest.java @@ -401,6 +401,27 @@ public void testConsumerClientRackSupplier() throws ConfigurationException, Clie assertThat(properties.get(PscConfiguration.PSC_CONSUMER_CLIENT_RACK)).isEqualTo(rackId); } + @Test + public void testFetchSideRateLimiterCreatedFromProps() throws ConfigurationException, ClientException { + Properties properties = new Properties(); + properties.setProperty( + com.pinterest.flink.connector.psc.source.PscSourceOptions.SCAN_RATE_LIMIT_RECORDS_PER_SECOND + .key(), + "10000"); + properties.setProperty(PscConfiguration.PSC_CONSUMER_POLL_MESSAGES_MAX, "250"); + PscTopicUriPartitionSplitReader reader = + createReader( + properties, UnregisteredMetricsGroup.createSourceReaderMetricGroup()); + assertThat(reader.fetchRateLimiter()).isNotNull(); + assertThat(reader.nextFetchRatePermits()).isEqualTo(250); + } + + @Test + public void testFetchSideRateLimiterAbsentWithoutProps() throws ConfigurationException, ClientException { + PscTopicUriPartitionSplitReader reader = createReader(); + assertThat(reader.fetchRateLimiter()).isNull(); + } + @ParameterizedTest @NullAndEmptySource public void testSetConsumerClientRackIgnoresNullAndEmpty(String rackId) throws ConfigurationException, ClientException { @@ -455,6 +476,7 @@ private void assignSplitsAndFetchUntilFinish(PscTopicUriPartitionSplitReader rea : recordCount + splitFetch.size()); splitId = recordsBySplitIds.nextSplit(); } + recordsBySplitIds.recycle(); } // Verify the number of records consumed from each split. diff --git a/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicTableFactoryTest.java b/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicTableFactoryTest.java index e747e49..6ca768f 100644 --- a/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicTableFactoryTest.java +++ b/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicTableFactoryTest.java @@ -1493,16 +1493,16 @@ public void testOperatorChainingWithRateLimitOnly() { // Get transformation using helper method (reuses same source) final Transformation transformation = produceTransformationFromSource(pscSource, 10); - // The transformation should be a OneInputTransformation (the rate limiter map) - // Since rescale is disabled, there should be no intermediate rescale transformation + // Rate limiting is fetch-side inside the source; no downstream PscRateLimit map. assertThat(transformation).isNotNull(); - assertThat(transformation.getName()).contains("PscRateLimit"); + assertThat(transformation).isInstanceOf(SourceTransformation.class); + assertThat(transformation.getName()).doesNotContain("PscRateLimit"); } @Test public void testOperatorChainingWithRescaleAndRateLimit() { // Verifies the new operator chain when both rescale and rate limiting are enabled: - // Source -> PscRateLimit -> rescale (terminal PartitionTransformation) + // Source -> rescale (terminal PartitionTransformation); rate limit is fetch-side // scan.parallelism (tier 1) determines the source/rate-limit parallelism, capped // by the env parallelism in produceTransformationFromSource. final Map modifiedOptions = @@ -1526,25 +1526,18 @@ public void testOperatorChainingWithRescaleAndRateLimit() { final Transformation terminal = produceTransformationFromSource(pscSource, envParallelism); - // Terminal is the rescale PartitionTransformation. + // Terminal is the rescale PartitionTransformation; rate limit is fetch-side (no map). assertThat(terminal).isNotNull(); assertThat(terminal).isInstanceOf(PartitionTransformation.class); - // Terminal's input is the rate-limit operator. assertThat(terminal.getInputs()).isNotEmpty(); - final Transformation rateLimitOp = terminal.getInputs().get(0); - assertThat(rateLimitOp.getName()).contains("PscRateLimit"); - - // The rate-limit operator wraps the Kafka source. - assertThat(rateLimitOp.getInputs()).isNotEmpty(); - final Transformation sourceOp = rateLimitOp.getInputs().get(0); + final Transformation sourceOp = terminal.getInputs().get(0); assertThat(sourceOp).isInstanceOf(SourceTransformation.class); + assertThat(sourceOp.getName()).doesNotContain("PscRateLimit"); - // Source and rate-limit are both pinned to min(scanParallelism, env). final int expectedSourceParallelism = Math.min(pscSource.scanParallelism, envParallelism); assertThat(sourceOp.getParallelism()).isEqualTo(expectedSourceParallelism); - assertThat(rateLimitOp.getParallelism()).isEqualTo(expectedSourceParallelism); } @Test @@ -1581,9 +1574,8 @@ public void testOperatorChainingWithRescaleWithoutRateLimit() { @Test public void testRateLimiterParallelismConfiguration() { // Verifies that scan.parallelism and rate-limit options are wired into the - // PscDynamicSource and that the rate-limit operator is inserted between the - // source and the rescale, with its parallelism matching the (capped) source - // parallelism. + // PscDynamicSource. Rate limiting is fetch-side; graph is Source -> rescale + // with source parallelism capped by env parallelism. final Map modifiedOptions = getModifiedOptions( getBasicSourceOptions(), @@ -1609,14 +1601,14 @@ public void testRateLimiterParallelismConfiguration() { assertThat(terminal).isInstanceOf(PartitionTransformation.class); assertThat(terminal.getInputs()).isNotEmpty(); - final Transformation rateLimitOp = terminal.getInputs().get(0); - assertThat(rateLimitOp.getName()).contains("PscRateLimit"); + // Rate limiting is fetch-side; rescale sits directly on the source. + final Transformation sourceOp = terminal.getInputs().get(0); + assertThat(sourceOp).isInstanceOf(SourceTransformation.class); + assertThat(sourceOp.getName()).doesNotContain("PscRateLimit"); - // Rate-limit parallelism follows the upstream source parallelism, which is - // capped at min(scanParallelism, env). final int expectedSourceParallelism = Math.min(pscSource.scanParallelism, envParallelism); - assertThat(rateLimitOp.getParallelism()).isEqualTo(expectedSourceParallelism); + assertThat(sourceOp.getParallelism()).isEqualTo(expectedSourceParallelism); } @Test @@ -1660,13 +1652,9 @@ public void testRescaleCreatesPartitionTransformation() { @Test public void testRescaleAndRateLimitChain() { - // Verifies the operator chain: Source -> RateLimit -> Rescale. - // Rescale is now applied AFTER the rate limiter, so the terminal transformation - // is a PartitionTransformation (the rescale), whose input is the rate-limit - // OneInputTransformation, whose input is the SourceTransformation. - // Also verifies that source/rate-limit parallelism is capped at env parallelism: + // Verifies the operator chain: Source -> Rescale (rate limit is fetch-side). + // Also verifies that source parallelism is capped at env parallelism: // sourceParallelism = min(scanParallelism, env.getParallelism()) = min(100, 10) = 10 - // rateLimiterParallelism = sourceStream.getParallelism() = 10 try { // Mock partition count = 20 (irrelevant to the gate now; rescale is gated only @@ -1705,23 +1693,14 @@ public void testRescaleAndRateLimitChain() { assertThat(terminal).isNotNull(); assertThat(terminal).isInstanceOf(PartitionTransformation.class); - // 2) Terminal's input is the rate-limit map. + // 2) Terminal's input is the source (rate limit is fetch-side). assertThat(terminal.getInputs()).isNotEmpty(); - final Transformation rateLimitOp = terminal.getInputs().get(0); - assertThat(rateLimitOp.getName()).contains("PscRateLimit"); + final Transformation sourceOp = terminal.getInputs().get(0); + assertThat(sourceOp).isInstanceOf(SourceTransformation.class); + assertThat(sourceOp.getName()).doesNotContain("PscRateLimit"); - // Rate-limit parallelism should match the (capped) source parallelism: - // min(scanParallelism=100, env=10) = 10. final int expectedSourceParallelism = Math.min(pscSource.scanParallelism, envParallelism); - assertThat(rateLimitOp.getParallelism()).isEqualTo(expectedSourceParallelism); - - // 3) The rate-limit map's input is the Kafka source. - assertThat(rateLimitOp.getInputs()).isNotEmpty(); - final Transformation sourceOp = rateLimitOp.getInputs().get(0); - assertThat(sourceOp).isInstanceOf(SourceTransformation.class); - - // Source operator parallelism should be capped at env parallelism. assertThat(sourceOp.getParallelism()).isEqualTo(expectedSourceParallelism); } finally { PscTableCommonUtils.resetProvider(); @@ -1783,11 +1762,10 @@ public void testRescaleAlwaysAppliedWhenEnabled() { @Test public void testRescaleAndRateLimitWithDifferentParallelism() { - // Verifies that when scan.parallelism is larger than env parallelism, both the - // source and the rate-limit operator are capped at env parallelism. The chain is: + // Verifies that when scan.parallelism is larger than env parallelism, the source + // is capped at env parallelism. Rate limit is fetch-side. The chain is: // Source(parallelism=min(scanParallelism, env)) - // -> PscRateLimit(parallelism=sourceStream.getParallelism()) - // -> rescale (terminal PartitionTransformation) + // -> rescale (terminal PartitionTransformation) try { // Mock partition count = 15. In the new logic this is not consulted because @@ -1821,20 +1799,14 @@ public void testRescaleAndRateLimitWithDifferentParallelism() { assertThat(terminal).isNotNull(); assertThat(terminal).isInstanceOf(PartitionTransformation.class); - // 2) Terminal's input is the rate-limit map. + // 2) Terminal's input is the source (rate limit is fetch-side). assertThat(terminal.getInputs()).isNotEmpty(); - final Transformation rateLimitOp = terminal.getInputs().get(0); - assertThat(rateLimitOp.getName()).contains("PscRateLimit"); + final Transformation sourceOp = terminal.getInputs().get(0); + assertThat(sourceOp).isInstanceOf(SourceTransformation.class); + assertThat(sourceOp.getName()).doesNotContain("PscRateLimit"); - // Rate-limit and source are both capped at env parallelism (NOT 80). final int expectedSourceParallelism = Math.min(pscSource.scanParallelism, envParallelism); - assertThat(rateLimitOp.getParallelism()).isEqualTo(expectedSourceParallelism); - - // 3) The rate-limit map's input is the Kafka source, also at the capped parallelism. - assertThat(rateLimitOp.getInputs()).isNotEmpty(); - final Transformation sourceOp = rateLimitOp.getInputs().get(0); - assertThat(sourceOp).isInstanceOf(SourceTransformation.class); assertThat(sourceOp.getParallelism()).isEqualTo(expectedSourceParallelism); } finally { PscTableCommonUtils.resetProvider(); @@ -2058,9 +2030,8 @@ public void testRescaleWithUnknownEffectiveParallelism() { @Test public void testRescaleAfterRateLimitOrder() { // Pins the operator ordering established in PscDynamicSource.produceDataStream(): - // Source -> PscRateLimit (rate limiter) -> rescale (terminal) - // i.e. rescale is the OUTERMOST operator, applied AFTER the rate limiter, - // not before it. This guards against accidental reordering. + // Source -> rescale (terminal); rate limit is applied fetch-side inside the source + // i.e. rescale is the OUTERMOST operator when enabled. final Map modifiedOptions = getModifiedOptions( getBasicSourceOptions(), @@ -2084,26 +2055,17 @@ public void testRescaleAfterRateLimitOrder() { assertThat(terminal).isInstanceOf(PartitionTransformation.class); assertThat(terminal.getName()).doesNotContain("PscRateLimit"); - // 2) Immediately below rescale is the rate limiter (NOT the bare source). + // 2) Immediately below rescale is the source (rate limit is fetch-side). assertThat(terminal.getInputs()).isNotEmpty(); - final Transformation rateLimitOp = terminal.getInputs().get(0); - assertThat(rateLimitOp.getName()).contains("PscRateLimit"); - assertThat(rateLimitOp).isNotInstanceOf(PartitionTransformation.class); - assertThat(rateLimitOp).isNotInstanceOf(SourceTransformation.class); - - // 3) Below the rate limiter is the Kafka source. This rules out the swapped - // ordering (Source -> rescale -> rate limiter), in which case the rate - // limiter's input would be a PartitionTransformation. - assertThat(rateLimitOp.getInputs()).isNotEmpty(); - final Transformation sourceOp = rateLimitOp.getInputs().get(0); + final Transformation sourceOp = terminal.getInputs().get(0); assertThat(sourceOp).isInstanceOf(SourceTransformation.class); + assertThat(sourceOp.getName()).doesNotContain("PscRateLimit"); } @Test public void testRateLimiterParallelismMatchesSourceParallelism() { - // The rate limiter is chained via DataStream.map(...) on the source stream, so its - // parallelism must equal the source stream's parallelism (the capped value), - // never scan.parallelism directly when scan.parallelism > env parallelism. + // Rate limiting is fetch-side inside the source. With rescale enabled, the graph is + // Source -> rescale, and source parallelism is capped at env parallelism. final Map modifiedOptions = getModifiedOptions( getBasicSourceOptions(), @@ -2124,20 +2086,12 @@ public void testRateLimiterParallelismMatchesSourceParallelism() { assertThat(terminal).isInstanceOf(PartitionTransformation.class); assertThat(terminal.getInputs()).isNotEmpty(); - final Transformation rateLimitOp = terminal.getInputs().get(0); - assertThat(rateLimitOp.getName()).contains("PscRateLimit"); - - assertThat(rateLimitOp.getInputs()).isNotEmpty(); - final Transformation sourceOp = rateLimitOp.getInputs().get(0); + final Transformation sourceOp = terminal.getInputs().get(0); assertThat(sourceOp).isInstanceOf(SourceTransformation.class); + assertThat(sourceOp.getName()).doesNotContain("PscRateLimit"); - // Both must equal min(scanParallelism, env). The rate limiter parallelism is - // explicitly read from sourceStream.getParallelism() in PscDynamicSource, so it - // is locked to the source operator's parallelism. final int expectedParallelism = Math.min(pscSource.scanParallelism, envParallelism); assertThat(sourceOp.getParallelism()).isEqualTo(expectedParallelism); - assertThat(rateLimitOp.getParallelism()).isEqualTo(expectedParallelism); - assertThat(rateLimitOp.getParallelism()).isEqualTo(sourceOp.getParallelism()); } @Test diff --git a/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/UpsertPscDynamicTableFactoryTest.java b/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/UpsertPscDynamicTableFactoryTest.java index 41ab741..4a5388c 100644 --- a/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/UpsertPscDynamicTableFactoryTest.java +++ b/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/UpsertPscDynamicTableFactoryTest.java @@ -753,7 +753,7 @@ public void testUpsertSourceWithRescaleEnabledProducesPartitionTransformation() @Test public void testUpsertSourceWithRescaleAndRateLimitOrder() { // Pins the operator ordering for the upsert source path: - // Source -> PscRateLimit -> rescale (terminal PartitionTransformation) + // Source -> rescale (terminal PartitionTransformation); rate limit is fetch-side // This must match the non-upsert factory's behavior. final Map options = getModifiedOptions( getFullSourceOptions(), @@ -779,21 +779,15 @@ public void testUpsertSourceWithRescaleAndRateLimitOrder() { assertThat(terminal).isInstanceOf(PartitionTransformation.class); assertThat(terminal.getName()).doesNotContain("PscRateLimit"); - // 2) Immediately below rescale is the rate limiter (NOT the bare source). + // 2) Immediately below rescale is the source (rate limit is fetch-side). assertThat(terminal.getInputs()).isNotEmpty(); - final Transformation rateLimitOp = terminal.getInputs().get(0); - assertThat(rateLimitOp.getName()).contains("PscRateLimit"); - - // 3) Below the rate limiter is the Kafka source. - assertThat(rateLimitOp.getInputs()).isNotEmpty(); - final Transformation sourceOp = rateLimitOp.getInputs().get(0); + final Transformation sourceOp = terminal.getInputs().get(0); assertThat(sourceOp).isInstanceOf(SourceTransformation.class); + assertThat(sourceOp.getName()).doesNotContain("PscRateLimit"); - // Source and rate-limit are both pinned to min(scanParallelism, env). final int expectedSourceParallelism = Math.min(pscSource.scanParallelism, envParallelism); assertThat(sourceOp.getParallelism()).isEqualTo(expectedSourceParallelism); - assertThat(rateLimitOp.getParallelism()).isEqualTo(expectedSourceParallelism); } @Test