Move source rate limiter to before poll()

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} <i>before</i> {@code consumer.poll()}, so MemQ/Kafka
+     * downloads are paced by the limiter (unlike a downstream map operator).
+     */
+    public static final ConfigOption<Double> 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> T getOption(
             Properties props, ConfigOption<?> configOption, Function<String, T> parser) {

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. " +

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<RowData> 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<RowData> 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();

psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscRateLimitMap.java

@@ -35,6 +35,10 @@
  * <p>When the rate limit is exceeded, the function blocks until permits become available,
  * emitting metrics to track throttling behavior.
  *
+ * <p><b>Note:</b> {@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 <T> The type of records flowing through this map function

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.