kafka实时消费优化

Pyspark_job/my_kafka/kafka_flow_asin_detail_to_doris.py

@@ -1150,23 +1150,32 @@ class KafkaFlowAsinDetail(Templates):
         end_time = time.time()
         print(f"Doris {self.doris_30day_table} 写入完毕，耗时：{end_time - start_time:.1f}s")
 
-    # 实时消费中批次数据的处理逻辑（latest 模式）
+    # 实时消费中批次数据的处理逻辑（latest 模式）：失败重试2次，仍失败则抛出异常（避免checkpoint误提交导致丢数据）
     def handle_kafka_stream(self, df, batch_id):
-        try:
-            batch_num = df.count()
-            if batch_num > 0:
+        max_retries = 3
+        for retry in range(max_retries):
+            try:
+                batch_num = df.count()
+                if batch_num == 0:
+                    print("当前批次没有数据")
+                    return
                 start_time = time.time()
                 print(f"当前批次:{batch_id}; 该批次数据量为：{batch_num}")
-                df = df.repartition(self.repartition_num)
-                df_save = self.handle_all_field(df)
+                df_repartitioned = df.repartition(self.repartition_num)
+                df_save = self.handle_all_field(df_repartitioned)
                 self.save_to_doris(df_save, batch_num)
                 df_save.unpersist()
                 end_time = time.time()
                 print(f"当前批次:{batch_id} 执行完毕, 执行时长为：{end_time - start_time:.1f}s")
-            else:
-                print("当前批次没有数据")
-        except Exception as e:
-            print(e, traceback.format_exc())
+                return
+            except Exception as e:
+                print(f"当前批次:{batch_id} 处理失败(第{retry + 1}次)，{e}", traceback.format_exc())
+                if retry == max_retries - 1:
+                    CommonUtil.send_wx_msg(
+                        ['chenyuanjie'], "⚠实时消费批次处理失败⚠", f"topic: {self.topic_name}, batch_id: {batch_id}, 重试{max_retries}次后仍失败: {e}"
+                    )
+                    raise
+                time.sleep(10)
 
     # 消费主题下的所有历史数据
     def handle_kafka_history(self, kafka_df):

Pyspark_job/utils/templates.py

@@ -349,6 +349,15 @@ class Templates(object):
             try:
                 if self.query is not None:
                     self.query.stop()  # 在子线程中调用，避免 foreachBatch 回调内死锁
+                    # stop()调用返回不代表引擎内部已彻底停止调度，显式轮询确认，
+                    # 避免后续补偿消费/spark.stop()和引擎残留调度撞车（IllegalStateException: stopped SparkContext）
+                    wait_seconds = 0
+                    while self.query.isActive and wait_seconds < 120:
+                        time.sleep(1)
+                        wait_seconds += 1
+                    if self.query.isActive:
+                        print(f"[停止] query.stop()后等待{wait_seconds}s仍处于active状态，继续执行后续流程")
+                self._compensate_missed_offsets()  # 停止前核对topic末尾offset，若有遗漏区间则补偿消费
                 if self.spark is not None:
                     self.spark.stop()
             except Exception as e:
@@ -360,6 +369,86 @@ class Templates(object):
         t.start()
         # foreachBatch 回调从此处正常返回，不阻塞等待 stop 完成
 
+    def _get_checkpoint_committed_offsets(self):
+        """解析checkpoint目录，取最后一次成功commit的批次对应的kafka分区offset；从未成功commit过则返回None"""
+        commit_files = HdfsUtils.read_list(f"{self.check_path}/commits")
+        if not commit_files:
+            return None
+        batch_ids = sorted(int(f) for f in commit_files if f.isdigit())
+        if not batch_ids:
+            return None
+        last_batch_id = batch_ids[-1]
+        offset_lines = HdfsUtils.read_hdfs_file(f"{self.check_path}/offsets/{last_batch_id}")
+        offset_json = json.loads(offset_lines[-1])  # 前面几行是版本号和元数据，最后一行才是offset
+        return offset_json.get(self.topic_name)
+
+    def _compensate_missed_offsets(self):
+        """停止消费前核对：对比kafka topic真实末尾offset与checkpoint最后一次成功commit的offset，
+        若有遗漏区间则以history方式分批（复用batch_size_history）补偿消费，确保topic消息全部被消费到"""
+        title_prefix = "[TEST] " if self.test_flag == 'test' else ""
+        try:
+            committed = self._get_checkpoint_committed_offsets()
+            consumer = self.get_kafka_object_by_python(topic_name=self.topic_name)
+            partition_data = self.get_kafka_partitions_data(consumer=consumer, topic_name=self.topic_name)
+            consumer.close()
+
+            current_offsets, target_end_offsets = {}, {}
+            gap_total = 0
+            for pid, info in partition_data.items():
+                end_offset = info['end_offsets']
+                begin_offset = info['beginning_offsets']
+                if committed is not None and str(pid) in committed:
+                    begin_offset = max(begin_offset, int(committed[str(pid)]))
+                if end_offset > begin_offset:
+                    current_offsets[pid] = begin_offset
+                    target_end_offsets[pid] = end_offset
+                    gap_total += end_offset - begin_offset
+
+            if gap_total == 0:
+                print(f"[停止前核对] {self.topic_name} 无遗漏，checkpoint offset已追平topic末尾")
+                return
+
+            print(f"[停止前核对] {self.topic_name} 发现遗漏共 {gap_total} 条，起始: {current_offsets} -> 目标: {target_end_offsets}，开始分批补偿消费")
+            compensated_total = 0
+            while current_offsets:
+                starting_offsets_dict, ending_offsets_dict = {}, {}
+                for pid, begin_offset in current_offsets.items():
+                    batch_end_offset = min(begin_offset + self.batch_size_history, target_end_offsets[pid])
+                    starting_offsets_dict[str(pid)] = begin_offset
+                    ending_offsets_dict[str(pid)] = batch_end_offset
+
+                starting_offsets_json = json.dumps({self.topic_name: starting_offsets_dict})
+                ending_offsets_json = json.dumps({self.topic_name: ending_offsets_dict})
+                kafka_df = self.create_kafka_df_object(
+                    consumer_type="history", topic_name=self.topic_name, schema=self.schema,
+                    starting_offsets_json=starting_offsets_json, ending_offsets_json=ending_offsets_json,
+                )
+                batch_num = kafka_df.count()
+                compensated_total += batch_num
+                print(f"[停止前核对] {self.topic_name} 补偿批次: {starting_offsets_dict} -> {ending_offsets_dict}，本批 {batch_num} 条")
+                self.handle_kafka_history(kafka_df)
+
+                for pid in list(current_offsets.keys()):
+                    new_offset = ending_offsets_dict[str(pid)]
+                    if new_offset >= target_end_offsets[pid]:
+                        del current_offsets[pid]
+                    else:
+                        current_offsets[pid] = new_offset
+
+            CommonUtil.send_wx_msg(
+                ['chenyuanjie'],
+                f"{title_prefix}实时消费停止前自动补偿完成: {self.topic_name}",
+                f"共补偿 {compensated_total} 条消息（差值估算 {gap_total} 条）"
+            )
+            print(f"[停止前核对] {self.topic_name} 补偿完成")
+        except Exception as e:
+            print(f"[停止前核对] {self.topic_name} 核对/补偿失败: {e}", traceback.format_exc())
+            CommonUtil.send_wx_msg(
+                ['chenyuanjie'],
+                f"{title_prefix}⚠实时消费停止前核对补偿失败⚠",
+                f"topic: {self.topic_name}，错误: {e}，需要人工核对kafka offset与Doris数据是否存在遗漏"
+            )
+
     def _start_state_monitor_thread(self, interval=900):
         """启动后台守护线程，独立轮询爬虫状态。
         解决 Kafka 无新数据时 forEachBatch 不触发、状态检查永远不执行的问题。"""
@@ -546,6 +635,7 @@ class Templates(object):
                         partition_offsets_dict[key]['beginning_offsets'] = value['beginning_offsets']
 
                 num = 0
+                next_beginning_offsets = {}  # 本批次成功处理后才生效的下一轮起始offset，避免失败重试时offset被提前跳过
                 for key, value in partition_offsets_dict.items():
                     # 起始偏移量
                     beginning_offsets = value['beginning_offsets']
@@ -557,7 +647,7 @@ class Templates(object):
                     if end_offsets >= end_offsets_partition:
                         num += 1
                     else:
-                        partition_offsets_dict[key]['beginning_offsets'] = end_offsets
+                        next_beginning_offsets[key] = end_offsets
 
                 starting_offsets_json = json.dumps({topic_name: beginning_offsets_dict})
                 ending_offsets_json = json.dumps({topic_name: end_offsets_dict})
@@ -577,14 +667,16 @@ class Templates(object):
                         continue
                 print(f"kafka_df.count():{kafka_df.count()}")
 
+                # 本批次真正处理成功（含记录offset）后，才把起始offset前移到下一轮，失败时下次重试仍从本批次起始offset开始
+                self.handle_kafka_history_templates(kafka_df=kafka_df)
+                self.record_offsets_by_history(end_offsets_dict=end_offsets_dict)
+                for key, new_begin in next_beginning_offsets.items():
+                    partition_offsets_dict[key]['beginning_offsets'] = new_begin
+
                 if num >= partition_num:
-                    self.handle_kafka_history_templates(kafka_df=kafka_df)  # 最后一批消费
-                    self.record_offsets_by_history(end_offsets_dict=end_offsets_dict)
                     self.start_process_instance()  # 退出之前启动调度
                     break
                 else:
-                    self.handle_kafka_history_templates(kafka_df=kafka_df)
-                    self.record_offsets_by_history(end_offsets_dict=end_offsets_dict)
                     time.sleep(10)
                     continue
 
@@ -595,7 +687,6 @@ class Templates(object):
 
     def handle_kafka_history_templates(self, kafka_df):
         self.handle_kafka_history(kafka_df)
-        self.kafka_consumption_is_finished()
 
     def handle_kafka_history(self, kafka_df):
         pass