Merge branch '1.10_release_4.0.x' into 1.10_release_4.1.x

Author: gituser

kafka-base/kafka-base-source/src/main/java/com/dtstack/flink/sql/source/kafka/KafkaDeserializationMetricWrapper.java

@@ -20,23 +20,24 @@
 
 import com.dtstack.flink.sql.dirtyManager.manager.DirtyDataManager;
 import com.dtstack.flink.sql.format.DeserializationMetricWrapper;
+import com.dtstack.flink.sql.util.ReflectionUtils;
 import org.apache.flink.api.common.serialization.DeserializationSchema;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.metrics.Gauge;
 import org.apache.flink.metrics.MetricGroup;
 import org.apache.flink.streaming.connectors.kafka.internal.KafkaConsumerThread;
 import org.apache.flink.streaming.connectors.kafka.internals.AbstractFetcher;
+import org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartition;
+import org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartitionState;
 import org.apache.flink.types.Row;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.consumer.internals.SubscriptionState;
 import org.apache.kafka.common.TopicPartition;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.IOException;
 import java.lang.reflect.Field;
-import java.util.Set;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.List;
 
 import static com.dtstack.flink.sql.metric.MetricConstant.DT_PARTITION_GROUP;
 import static com.dtstack.flink.sql.metric.MetricConstant.DT_TOPIC_GROUP;
@@ -53,10 +54,6 @@ public class KafkaDeserializationMetricWrapper extends DeserializationMetricWrap
 
     private static final Logger LOG = LoggerFactory.getLogger(KafkaDeserializationMetricWrapper.class);
 
-    private AbstractFetcher<Row, ?> fetcher;
-
-    private AtomicBoolean firstMsg = new AtomicBoolean(true);
-
     private Calculate calculate;
 
     public KafkaDeserializationMetricWrapper(
@@ -68,66 +65,82 @@ public KafkaDeserializationMetricWrapper(
         this.calculate = calculate;
     }
 
-    @Override
-    protected void beforeDeserialize() throws IOException {
-        super.beforeDeserialize();
-        if (firstMsg.compareAndSet(true, false)) {
-            try {
-                registerPtMetric(fetcher);
-            } catch (Exception e) {
-                LOG.error("register topic partition metric error.", e);
-            }
-        }
-    }
-
     protected void registerPtMetric(AbstractFetcher<Row, ?> fetcher) throws Exception {
-        Field consumerThreadField = getConsumerThreadField(fetcher);
+        Field consumerThreadField = ReflectionUtils.getDeclaredField(fetcher, "consumerThread");
         consumerThreadField.setAccessible(true);
         KafkaConsumerThread consumerThread = (KafkaConsumerThread) consumerThreadField.get(fetcher);
 
         Field hasAssignedPartitionsField = consumerThread.getClass().getDeclaredField("hasAssignedPartitions");
         hasAssignedPartitionsField.setAccessible(true);
 
-        //wait until assignedPartitions
-
-        boolean hasAssignedPartitions = (boolean) hasAssignedPartitionsField.get(consumerThread);
+        // get subtask unassigned kafka topic partition
+        Field subscribedPartitionStatesField = ReflectionUtils.getDeclaredField(fetcher, "subscribedPartitionStates");
+        subscribedPartitionStatesField.setAccessible(true);
+        List<KafkaTopicPartitionState<KafkaTopicPartition>> subscribedPartitionStates = (List<KafkaTopicPartitionState<KafkaTopicPartition>>) subscribedPartitionStatesField.get(fetcher);
 
-        if (!hasAssignedPartitions) {
-            throw new RuntimeException("wait 50 secs, but not assignedPartitions");
-        }
+        // init partition lag metric
+        for (KafkaTopicPartitionState<KafkaTopicPartition> kafkaTopicPartitionState : subscribedPartitionStates) {
+            KafkaTopicPartition kafkaTopicPartition = kafkaTopicPartitionState.getKafkaTopicPartition();
+            MetricGroup topicMetricGroup = getRuntimeContext().getMetricGroup().addGroup(DT_TOPIC_GROUP, kafkaTopicPartition.getTopic());
 
-        Field consumerField = consumerThread.getClass().getDeclaredField("consumer");
-        consumerField.setAccessible(true);
-
-        KafkaConsumer kafkaConsumer = (KafkaConsumer) consumerField.get(consumerThread);
-        Field subscriptionStateField = kafkaConsumer.getClass().getDeclaredField("subscriptions");
-        subscriptionStateField.setAccessible(true);
-
-        //topic partitions lag
-        SubscriptionState subscriptionState = (SubscriptionState) subscriptionStateField.get(kafkaConsumer);
-        Set<TopicPartition> assignedPartitions = subscriptionState.assignedPartitions();
-
-        for (TopicPartition topicPartition : assignedPartitions) {
-            MetricGroup metricGroup = getRuntimeContext().getMetricGroup().addGroup(DT_TOPIC_GROUP, topicPartition.topic())
-                    .addGroup(DT_PARTITION_GROUP, topicPartition.partition() + "");
+            MetricGroup metricGroup = topicMetricGroup.addGroup(DT_PARTITION_GROUP, kafkaTopicPartition.getPartition() + "");
             metricGroup.gauge(DT_TOPIC_PARTITION_LAG_GAUGE, new Gauge<Long>() {
+                // tmp variable
+                boolean initLag = true;
+                int partitionIndex;
+                SubscriptionState subscriptionState;
+                TopicPartition topicPartition;
+
                 @Override
                 public Long getValue() {
-                    return calculate.calc(subscriptionState, topicPartition);
+                    // first time register metrics
+                    if (initLag) {
+                        partitionIndex = kafkaTopicPartition.getPartition();
+                        initLag = false;
+                        return -1L;
+                    }
+                    // when kafka topic partition assigned calc metrics
+                    if (subscriptionState == null) {
+                        try {
+                            Field consumerField = consumerThread.getClass().getDeclaredField("consumer");
+                            consumerField.setAccessible(true);
+
+                            KafkaConsumer kafkaConsumer = (KafkaConsumer) consumerField.get(consumerThread);
+                            Field subscriptionStateField = kafkaConsumer.getClass().getDeclaredField("subscriptions");
+                            subscriptionStateField.setAccessible(true);
+
+                            boolean hasAssignedPartitions = (boolean) hasAssignedPartitionsField.get(consumerThread);
+
+                            if (!hasAssignedPartitions) {
+                                LOG.error("wait 50 secs, but not assignedPartitions");
+                            }
+
+                            subscriptionState = (SubscriptionState) subscriptionStateField.get(kafkaConsumer);
+
+                            topicPartition = subscriptionState
+                                    .assignedPartitions()
+                                    .stream()
+                                    .filter(x -> x.partition() == partitionIndex)
+                                    .findFirst()
+                                    .get();
+
+                        } catch (Exception e) {
+                            LOG.error(e.getMessage());
+                        }
+                        return -1L;
+                    } else {
+                        return calculate.calc(subscriptionState, topicPartition);
+                    }
                 }
             });
         }
     }
 
     public void setFetcher(AbstractFetcher<Row, ?> fetcher) {
-        this.fetcher = fetcher;
-    }
-
-    private Field getConsumerThreadField(AbstractFetcher fetcher) throws NoSuchFieldException {
         try {
-            return fetcher.getClass().getDeclaredField("consumerThread");
+            registerPtMetric(fetcher);
         } catch (Exception e) {
-            return fetcher.getClass().getSuperclass().getDeclaredField("consumerThread");
+            LOG.error("register topic partition metric error.", e);
         }
     }
 }