[INLONG-9380][Sort] Audit lost when stop job immediately after checkpoint

inlong-sort/sort-flink/base/src/main/java/org/apache/inlong/sort/base/metric/SourceMetricData.java

@@ -298,16 +298,6 @@ public void outputMetrics(long rowCountSize, long rowDataSize, long fetchDelay,
         }
     }
 
-    /**
-     * flush audit data
-     * usually call this method in close method or when checkpointing
-     */
-    public void flushAuditData() {
-        if (auditOperator != null) {
-            auditOperator.send();
-        }
-    }
-
     public void outputMetrics(long rowCountSize, long rowDataSize, long dataTime) {
         outputDefaultMetrics(rowCountSize, rowDataSize);
         if (auditOperator != null) {
@@ -345,6 +335,16 @@ private void outputDefaultMetrics(long rowCountSize, long rowDataSize) {
         }
     }
 
+    /**
+     * flush audit data
+     * usually call this method in close method or when checkpointing
+     */
+    public void flushAuditData() {
+        if (auditOperator != null) {
+            auditOperator.send();
+        }
+    }
+
     private void outputDefaultMetrics(long rowCountSize, long rowDataSize, long fetchDelay, long emitDelay) {
         outputDefaultMetrics(rowCountSize, rowDataSize);
         this.fetchDelay = fetchDelay;

inlong-sort/sort-flink/sort-flink-v1.15/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/IcebergStreamWriter.java

@@ -20,6 +20,7 @@
 import org.apache.inlong.sort.base.metric.MetricOption;
 import org.apache.inlong.sort.iceberg.utils.SinkMetadataUtils;
 
+import org.apache.flink.runtime.state.StateSnapshotContext;
 import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
 import org.apache.flink.streaming.api.operators.BoundedOneInput;
 import org.apache.flink.streaming.api.operators.ChainingStrategy;
@@ -78,6 +79,11 @@ public void prepareSnapshotPreBarrier(long checkpointId) throws Exception {
         this.writer = taskWriterFactory.create();
     }
 
+    @Override
+    public void snapshotState(StateSnapshotContext context) {
+        writerMetrics.flushAudit();
+    }
+
     @Override
     public void processElement(StreamRecord<T> element) throws Exception {
         T data = element.getValue();

inlong-sort/sort-flink/sort-flink-v1.15/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/sink/IcebergStreamWriterMetrics.java

@@ -113,4 +113,10 @@ void outputMetricsWithEstimate(int size, long time) {
             sourceMetricData.outputMetrics(1, size, time);
         }
     }
+
+    void flushAudit() {
+        if (sourceMetricData != null) {
+            sourceMetricData.flushAuditData();
+        }
+    }
 }

inlong-sort/sort-flink/sort-flink-v1.15/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/source/reader/IcebergSourceReader.java

@@ -28,6 +28,7 @@
 
 import java.util.Collection;
 import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -38,6 +39,7 @@ public class IcebergSourceReader<T>
         extends
             SingleThreadMultiplexSourceReaderBase<RecordAndPosition<T>, T, IcebergSourceSplit, IcebergSourceSplit> {
 
+    private final InlongIcebergSourceReaderMetrics<T> metrics;
     public IcebergSourceReader(
             InlongIcebergSourceReaderMetrics<T> metrics,
             ReaderFunction<T> readerFunction,
@@ -47,6 +49,7 @@ public IcebergSourceReader(
                 new IcebergSourceRecordEmitter<>(),
                 context.getConfiguration(),
                 context);
+        this.metrics = metrics;
     }
 
     @Override
@@ -62,6 +65,11 @@ public void start() {
     protected void onSplitFinished(Map<String, IcebergSourceSplit> finishedSplitIds) {
         requestSplit(Lists.newArrayList(finishedSplitIds.keySet()));
     }
+    @Override
+    public List<IcebergSourceSplit> snapshotState(long checkpointId) {
+        metrics.flushAudit();
+        return super.snapshotState(checkpointId);
+    }
 
     @Override
     protected IcebergSourceSplit initializedState(IcebergSourceSplit split) {

inlong-sort/sort-flink/sort-flink-v1.15/sort-connectors/iceberg/src/main/java/org/apache/inlong/sort/iceberg/source/reader/InlongIcebergSourceReaderMetrics.java

@@ -77,4 +77,10 @@ private long getDataSize(T object) {
         }
         return object.toString().getBytes(StandardCharsets.UTF_8).length;
     }
+
+    void flushAudit() {
+        if (sourceMetricData != null) {
+            sourceMetricData.flushAuditData();
+        }
+    }
 }

inlong-sort/sort-flink/sort-flink-v1.15/sort-connectors/tubemq/src/main/java/org/apache/inlong/sort/tubemq/FlinkTubeMQConsumer.java

@@ -317,6 +317,8 @@ public void snapshotState(FunctionSnapshotContext context) throws Exception {
             offsetsState.add(new Tuple2<>(entry.getKey(), entry.getValue()));
         }
 
+        deserializationSchema.flushAudit();
+
         LOG.info("Successfully save the offsets in checkpoint {}: {}.",
                 context.getCheckpointId(), currentOffsets);
     }

inlong-sort/sort-flink/sort-flink-v1.15/sort-connectors/tubemq/src/main/java/org/apache/inlong/sort/tubemq/table/DynamicTubeMQDeserializationSchema.java

@@ -58,4 +58,6 @@ default void deserialize(Message message, Collector<T> out) throws IOException {
             out.collect(deserialize);
         }
     }
+
+    void flushAudit();
 }

inlong-sort/sort-flink/sort-flink-v1.15/sort-connectors/tubemq/src/main/java/org/apache/inlong/sort/tubemq/table/DynamicTubeMQTableDeserializationSchema.java

@@ -109,6 +109,13 @@ public void deserialize(Message message, Collector<RowData> out) throws IOExcept
 
     }
 
+    @Override
+    public void flushAudit() {
+        if (sourceMetricData != null) {
+            sourceMetricData.flushAuditData();
+        }
+    }
+
     @Override
     public TypeInformation<RowData> getProducedType() {
         return producedTypeInfo;