[FLINK-29398][connector/kafka] Provide rack ID to Kafka Source to take advantage of Rack Awareness

docs/content/docs/connectors/datastream/kafka.md

@@ -466,6 +466,28 @@ client dependencies in the job JAR, so you may need to rewrite it with the actua
 For detailed explanations of security configurations, please refer to
 <a href="https://kafka.apache.org/documentation/#security">the "Security" section in Apache Kafka documentation</a>.
 
+## Kafka Rack Awareness
+
+Kafka Rack Awareness allows flink to select and control the cloud region and availability zone configured by the use of rackId, this feature could allow a significant cost reduction in the cloud provider bill and achieve a better networking performance when connecting to closer and more reliable networks.
+
+### RackId
+
+setRackId is the variable where the desired or available availability zones get stored. If provided, the Supplier will be run when the consumer is set up on the Task Manager, and the consumer's client.rack configuration will be set to the value.
+
+https://kafka.apache.org/documentation/#consumerconfigs_client.rack
+
+One of the ways this can be implemented is by making setRackId equal to an environment variable within your taskManager, for instance:
+
+```
+.setRackId(() -> System.getenv("TM_NODE_AZ"))
+```
+
+The "TM_NODE_AZ" is the name of the environment variable in the TaskManager image that contains the available Network zones we want to use. 
+
+Another implementation option could be extracting the AZ directly from AWS CLI or API.
+
+Additional validation is added to ensure any input to the supplier gets properly configured by the consumer and ensures null values are handled properly.
+
 ### Behind the Scene
 {{< hint info >}}
 If you are interested in how Kafka source works under the design of new data source API, you may

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/KafkaSource.java

@@ -98,20 +98,24 @@ public class KafkaSource<OUT>
     private final KafkaRecordDeserializationSchema<OUT> deserializationSchema;
     // The configurations.
     private final Properties props;
+    // Client rackId callback
+    private final Supplier<String> rackIdSupplier;
 
     KafkaSource(
             KafkaSubscriber subscriber,
             OffsetsInitializer startingOffsetsInitializer,
             @Nullable OffsetsInitializer stoppingOffsetsInitializer,
             Boundedness boundedness,
             KafkaRecordDeserializationSchema<OUT> deserializationSchema,
-            Properties props) {
+            Properties props,
+            Supplier<String> rackIdSupplier) {
         this.subscriber = subscriber;
         this.startingOffsetsInitializer = startingOffsetsInitializer;
         this.stoppingOffsetsInitializer = stoppingOffsetsInitializer;
         this.boundedness = boundedness;
         this.deserializationSchema = deserializationSchema;
         this.props = props;
+        this.rackIdSupplier = rackIdSupplier;
     }
 
     /**
@@ -157,7 +161,9 @@ public UserCodeClassLoader getUserCodeClassLoader() {
                 new KafkaSourceReaderMetrics(readerContext.metricGroup());
 
         Supplier<KafkaPartitionSplitReader> splitReaderSupplier =
-                () -> new KafkaPartitionSplitReader(props, readerContext, kafkaSourceReaderMetrics);
+                () ->
+                        new KafkaPartitionSplitReader(
+                                props, readerContext, kafkaSourceReaderMetrics, rackIdSupplier);
         KafkaRecordEmitter<OUT> recordEmitter = new KafkaRecordEmitter<>(deserializationSchema);
 
         return new KafkaSourceReader<>(

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/KafkaSourceBuilder.java

@@ -40,6 +40,7 @@
 import java.util.Properties;
 import java.util.Random;
 import java.util.Set;
+import java.util.function.Supplier;
 import java.util.regex.Pattern;
 
 import static org.apache.flink.util.Preconditions.checkNotNull;
@@ -80,6 +81,7 @@
  *     .setTopics(Arrays.asList(TOPIC1, TOPIC2))
  *     .setDeserializer(KafkaRecordDeserializationSchema.valueOnly(StringDeserializer.class))
  *     .setUnbounded(OffsetsInitializer.latest())
+ *     .setRackId(() -> MY_RACK_ID)
  *     .build();
  * }</pre>
  *
@@ -100,6 +102,8 @@ public class KafkaSourceBuilder<OUT> {
     private KafkaRecordDeserializationSchema<OUT> deserializationSchema;
     // The configurations.
     protected Properties props;
+    // Client rackId supplier
+    private Supplier<String> rackIdSupplier;
 
     KafkaSourceBuilder() {
         this.subscriber = null;
@@ -108,6 +112,7 @@ public class KafkaSourceBuilder<OUT> {
         this.boundedness = Boundedness.CONTINUOUS_UNBOUNDED;
         this.deserializationSchema = null;
         this.props = new Properties();
+        this.rackIdSupplier = null;
     }
 
     /**
@@ -355,6 +360,17 @@ public KafkaSourceBuilder<OUT> setClientIdPrefix(String prefix) {
         return setProperty(KafkaSourceOptions.CLIENT_ID_PREFIX.key(), prefix);
     }
 
+    /**
+     * Set the clientRackId supplier to be passed down to the KafkaPartitionSplitReader.
+     *
+     * @param rackIdCallback callback to provide Kafka consumer client.rack
+     * @return this KafkaSourceBuilder
+     */
+    public KafkaSourceBuilder<OUT> setRackId(Supplier<String> rackIdCallback) {
+        this.rackIdSupplier = rackIdCallback;
+        return this;
+    }
+
     /**
      * Set an arbitrary property for the KafkaSource and KafkaConsumer. The valid keys can be found
      * in {@link ConsumerConfig} and {@link KafkaSourceOptions}.
@@ -422,7 +438,8 @@ public KafkaSource<OUT> build() {
                 stoppingOffsetsInitializer,
                 boundedness,
                 deserializationSchema,
-                props);
+                props,
+                rackIdSupplier);
     }
 
     // ------------- private helpers  --------------

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/KafkaPartitionSplitReader.java

@@ -80,11 +80,20 @@ public KafkaPartitionSplitReader(
             Properties props,
             SourceReaderContext context,
             KafkaSourceReaderMetrics kafkaSourceReaderMetrics) {
+        this(props, context, kafkaSourceReaderMetrics, () -> null);
+    }
+
+    public KafkaPartitionSplitReader(
+            Properties props,
+            SourceReaderContext context,
+            KafkaSourceReaderMetrics kafkaSourceReaderMetrics,
+            Supplier<String> rackIdSupplier) {
         this.subtaskId = context.getIndexOfSubtask();
         this.kafkaSourceReaderMetrics = kafkaSourceReaderMetrics;
         Properties consumerProps = new Properties();
         consumerProps.putAll(props);
         consumerProps.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, createConsumerClientId(props));
+        setConsumerClientRack(consumerProps, rackIdSupplier);
         this.consumer = new KafkaConsumer<>(consumerProps);
         this.stoppingOffsets = new HashMap<>();
         this.groupId = consumerProps.getProperty(ConsumerConfig.GROUP_ID_CONFIG);
@@ -256,6 +265,22 @@ KafkaConsumer<byte[], byte[]> consumer() {
 
     // --------------- private helper method ----------------------
 
+    /**
+     * This Method performs Null and empty Rack Id validation and sets the rack id to the
+     * client.rack Consumer Config.
+     *
+     * @param consumerProps Consumer Property.
+     * @param rackIdSupplier Rack Id's.
+     */
+    void setConsumerClientRack(Properties consumerProps, Supplier<String> rackIdSupplier) {
+        if (rackIdSupplier != null) {
+            String rackId = rackIdSupplier.get();
+            if (rackId != null && !rackId.isEmpty()) {
+                consumerProps.setProperty(ConsumerConfig.CLIENT_RACK_CONFIG, rackId);
+            }
+        }
+    }
+
     private void parseStartingOffsets(
             KafkaPartitionSplit split,
             List<TopicPartition> partitionsStartingFromEarliest,

flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/reader/KafkaPartitionSplitReaderTest.java

@@ -48,6 +48,7 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.CsvSource;
 import org.junit.jupiter.params.provider.EmptySource;
+import org.junit.jupiter.params.provider.NullAndEmptySource;
 import org.junit.jupiter.params.provider.ValueSource;
 
 import java.io.IOException;
@@ -63,6 +64,7 @@
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 
 import static org.apache.flink.connector.kafka.testutils.KafkaSourceTestEnv.NUM_RECORDS_PER_PARTITION;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -319,6 +321,51 @@ public void testUsingCommittedOffsetsWithEarliestOrLatestOffsetResetStrategy(
         assertThat(reader.consumer().position(partition)).isEqualTo(expectedOffset);
     }
 
+    @Test
+    public void testConsumerClientRackSupplier() {
+        AtomicReference<Boolean> supplierCalled = new AtomicReference<>(false);
+        String rackId = "use1-az1";
+        Supplier<String> rackIdSupplier =
+                () -> {
+                    supplierCalled.set(true);
+                    return rackId;
+                };
+        Properties properties = new Properties();
+        createReader(
+                properties,
+                UnregisteredMetricsGroup.createSourceReaderMetricGroup(),
+                rackIdSupplier);
+        assertThat(supplierCalled.get()).isEqualTo(true);
+    }
+
+    @ParameterizedTest
+    @NullAndEmptySource
+    public void testSetConsumerClientRackIgnoresNullAndEmpty(String rackId) {
+        Properties properties = new Properties();
+        Supplier<String> rackIdSupplier = () -> rackId;
+        KafkaPartitionSplitReader reader =
+                createReader(
+                        properties,
+                        UnregisteredMetricsGroup.createSourceReaderMetricGroup(),
+                        rackIdSupplier);
+        reader.setConsumerClientRack(properties, rackIdSupplier);
+        assertThat(properties.containsKey(ConsumerConfig.CLIENT_RACK_CONFIG)).isFalse();
+    }
+
+    @Test
+    public void testSetConsumerClientRackUsesCorrectParameter() {
+        String rackId = "use1-az1";
+        Properties properties = new Properties();
+        Supplier<String> rackIdSupplier = () -> rackId;
+        KafkaPartitionSplitReader reader =
+                createReader(
+                        properties,
+                        UnregisteredMetricsGroup.createSourceReaderMetricGroup(),
+                        rackIdSupplier);
+        reader.setConsumerClientRack(properties, rackIdSupplier);
+        assertThat(properties.get(ConsumerConfig.CLIENT_RACK_CONFIG)).isEqualTo(rackId);
+    }
+
     // ------------------
 
     private void assignSplitsAndFetchUntilFinish(KafkaPartitionSplitReader reader, int readerId)
@@ -383,6 +430,13 @@ private KafkaPartitionSplitReader createReader() {
 
     private KafkaPartitionSplitReader createReader(
             Properties additionalProperties, SourceReaderMetricGroup sourceReaderMetricGroup) {
+        return createReader(additionalProperties, sourceReaderMetricGroup, () -> null);
+    }
+
+    private KafkaPartitionSplitReader createReader(
+            Properties additionalProperties,
+            SourceReaderMetricGroup sourceReaderMetricGroup,
+            Supplier<String> rackIdSupplier) {
         Properties props = new Properties();
         props.putAll(KafkaSourceTestEnv.getConsumerProperties(ByteArrayDeserializer.class));
         props.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "none");
@@ -394,7 +448,8 @@ private KafkaPartitionSplitReader createReader(
         return new KafkaPartitionSplitReader(
                 props,
                 new TestingReaderContext(new Configuration(), sourceReaderMetricGroup),
-                kafkaSourceReaderMetrics);
+                kafkaSourceReaderMetrics,
+                rackIdSupplier);
     }
 
     private Map<String, KafkaPartitionSplit> assignSplits(