Pipe: Modify epoch status metric changes (#16272)

* Pipe: Modify epoch status metric changes

* fix

Author: luoluoyuyu

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/source/dataregion/realtime/PipeRealtimeDataRegionHybridSource.java

@@ -53,6 +53,7 @@ protected void doExtract(final PipeRealtimeEvent event) {
       extractTabletInsertion(event);
     } else if (eventToExtract instanceof TsFileInsertionEvent) {
       extractTsFileInsertion(event);
+      event.getTsFileEpoch().clearState(this);
     } else if (eventToExtract instanceof PipeHeartbeatEvent) {
       extractHeartbeat(event);
     } else if (eventToExtract instanceof PipeDeleteDataNodeEvent) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/source/dataregion/realtime/PipeRealtimeDataRegionLogSource.java

@@ -47,6 +47,7 @@ protected void doExtract(PipeRealtimeEvent event) {
       extractTabletInsertion(event);
     } else if (eventToExtract instanceof TsFileInsertionEvent) {
       extractTsFileInsertion(event);
+      event.getTsFileEpoch().clearState(this);
     } else if (eventToExtract instanceof PipeHeartbeatEvent) {
       extractHeartbeat(event);
     } else if (eventToExtract instanceof PipeDeleteDataNodeEvent) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/source/dataregion/realtime/PipeRealtimeDataRegionSource.java

@@ -63,6 +63,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
@@ -121,6 +122,8 @@ public abstract class PipeRealtimeDataRegionSource implements PipeExtractor {
   private boolean sloppyTimeRange; // true to disable time range filter after extraction
   private boolean sloppyPattern; // true to disable pattern filter after extraction
 
+  private AtomicLong extractEpochSize = new AtomicLong();
+
   // This queue is used to store pending events extracted by the method extract(). The method
   // supply() will poll events from this queue and send them to the next pipe plugin.
   protected final UnboundedBlockingPendingQueue<Event> pendingQueue =
@@ -646,4 +649,16 @@ public int getEventCount() {
   public String getTaskID() {
     return taskID;
   }
+
+  public void increaseExtractEpochSize() {
+    extractEpochSize.incrementAndGet();
+  }
+
+  public void decreaseExtractEpochSize() {
+    extractEpochSize.decrementAndGet();
+  }
+
+  public boolean extractEpochSizeIsEmpty() {
+    return extractEpochSize.get() == 0;
+  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/source/dataregion/realtime/PipeRealtimeDataRegionTsFileSource.java

@@ -74,6 +74,8 @@ protected void doExtract(PipeRealtimeEvent event) {
       // Ignore the event.
       event.decreaseReferenceCount(PipeRealtimeDataRegionTsFileSource.class.getName(), false);
     }
+
+    event.getTsFileEpoch().clearState(this);
   }
 
   @Override

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/source/dataregion/realtime/epoch/TsFileEpoch.java

@@ -23,6 +23,7 @@
 import org.apache.iotdb.db.pipe.source.dataregion.realtime.PipeRealtimeDataRegionSource;
 import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource;
 
+import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicLong;
@@ -49,9 +50,30 @@ public TsFileEpoch.State getState(final PipeRealtimeDataRegionSource extractor)
 
   public void migrateState(
       final PipeRealtimeDataRegionSource extractor, final TsFileEpochStateMigrator visitor) {
-    dataRegionExtractor2State
-        .computeIfAbsent(extractor, o -> new AtomicReference<>(State.EMPTY))
-        .getAndUpdate(visitor::migrate);
+    AtomicReference<State> stateRef = dataRegionExtractor2State.get(extractor);
+
+    if (stateRef == null) {
+      dataRegionExtractor2State.putIfAbsent(
+          extractor, stateRef = new AtomicReference<>(State.EMPTY));
+      extractor.increaseExtractEpochSize();
+      setExtractorsRecentProcessedTsFileEpochState();
+    }
+
+    State migratedState = visitor.migrate(stateRef.get());
+    if (!Objects.equals(stateRef.get(), migratedState)) {
+      stateRef.set(migratedState);
+      setExtractorsRecentProcessedTsFileEpochState();
+    }
+  }
+
+  public void clearState(final PipeRealtimeDataRegionSource extractor) {
+    if (dataRegionExtractor2State.containsKey(extractor)) {
+      extractor.decreaseExtractEpochSize();
+    }
+    if (extractor.extractEpochSizeIsEmpty()) {
+      PipeDataRegionSourceMetrics.getInstance()
+          .setRecentProcessedTsFileEpochState(extractor.getTaskID(), State.EMPTY);
+    }
   }
 
   public void setExtractorsRecentProcessedTsFileEpochState() {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/source/dataregion/realtime/epoch/TsFileEpochManager.java

@@ -61,9 +61,6 @@ public PipeRealtimeEvent bindPipeTsFileInsertionEvent(
         });
 
     final TsFileEpoch epoch = filePath2Epoch.remove(filePath);
-    // When all data corresponding to this TsFileEpoch have been extracted, update the state
-    // of the extractors processing this TsFileEpoch.
-    epoch.setExtractorsRecentProcessedTsFileEpochState();
 
     LOGGER.info("All data in TsFileEpoch {} was extracted", epoch);
     return new PipeRealtimeEvent(