[Backport 1.3] Fix write retries, Pipe resource cleanup, and snapshot recovery

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigNodeRuntimeAgent.java

@@ -24,6 +24,7 @@
 import org.apache.iotdb.commons.log.LoggerPeriodicalLogReducer;
 import org.apache.iotdb.commons.pipe.agent.runtime.PipePeriodicalJobExecutor;
 import org.apache.iotdb.commons.pipe.agent.runtime.PipePeriodicalPhantomReferenceCleaner;
+import org.apache.iotdb.commons.pipe.agent.task.meta.PipeMeta;
 import org.apache.iotdb.commons.pipe.agent.task.meta.PipeTaskMeta;
 import org.apache.iotdb.commons.pipe.config.PipeConfig;
 import org.apache.iotdb.commons.pipe.event.EnrichedEvent;
@@ -126,6 +127,11 @@ public void decreaseListenerReference(final PipeParameters parameters)
     regionListener.decreaseReference(parameters);
   }
 
+  public void reconcileListenerReferences(final Iterable<PipeMeta> pipeMetas)
+      throws IllegalPathException {
+    regionListener.reconcileReferences(pipeMetas);
+  }
+
   /** Notify the region listener that the leader is ready to allow pipe operations. */
   public void notifyLeaderReady() {
     regionListener.notifyLeaderReady();

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigRegionListener.java

@@ -20,6 +20,7 @@
 package org.apache.iotdb.confignode.manager.pipe.agent.runtime;
 
 import org.apache.iotdb.commons.exception.IllegalPathException;
+import org.apache.iotdb.commons.pipe.agent.task.meta.PipeMeta;
 import org.apache.iotdb.confignode.manager.pipe.source.ConfigRegionListeningFilter;
 import org.apache.iotdb.confignode.manager.pipe.source.ConfigRegionListeningQueue;
 import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters;
@@ -57,6 +58,24 @@ public synchronized void decreaseReference(final PipeParameters parameters)
     }
   }
 
+  public synchronized void reconcileReferences(final Iterable<PipeMeta> pipeMetas)
+      throws IllegalPathException {
+    int referenceCount = 0;
+    for (final PipeMeta pipeMeta : pipeMetas) {
+      if (!ConfigRegionListeningFilter.parseListeningPlanTypeSet(
+              pipeMeta.getStaticMeta().getExtractorParameters())
+          .isEmpty()) {
+        referenceCount++;
+      }
+    }
+    listeningQueueReferenceCount = referenceCount;
+    if (referenceCount == 0) {
+      listeningQueue.close();
+    } else {
+      listeningQueue.open();
+    }
+  }
+
   public synchronized boolean isLeaderReady() {
     return isLeaderReady.get();
   }

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java

@@ -22,6 +22,7 @@
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.common.rpc.thrift.TSchemaNode;
 import org.apache.iotdb.commons.auth.AuthException;
+import org.apache.iotdb.commons.exception.IllegalPathException;
 import org.apache.iotdb.commons.path.PartialPath;
 import org.apache.iotdb.commons.schema.node.MNodeType;
 import org.apache.iotdb.commons.schema.ttl.TTLCache;
@@ -663,9 +664,16 @@ public void loadSnapshot(File latestSnapshotRootDir) {
               }
             });
     if (result.get()) {
-      LOGGER.info(
-          "[ConfigNodeSnapshot] Load snapshot success, latestSnapshotRootDir: {}",
-          latestSnapshotRootDir);
+      try {
+        PipeConfigNodeAgent.runtime()
+            .reconcileListenerReferences(pipeInfo.getPipeTaskInfo().getPipeMetaList());
+        LOGGER.info(
+            "[ConfigNodeSnapshot] Load snapshot success, latestSnapshotRootDir: {}",
+            latestSnapshotRootDir);
+      } catch (final IllegalPathException e) {
+        result.set(false);
+        LOGGER.error(e.getMessage(), e);
+      }
     }
   }
 

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/pipe/PipeInfo.java

@@ -290,11 +290,6 @@ public void processLoadSnapshot(final File snapshotDir) throws IOException {
 
     try {
       pipeTaskInfo.processLoadSnapshot(snapshotDir);
-
-      for (final PipeMeta pipeMeta : pipeTaskInfo.getPipeMetaList()) {
-        PipeConfigNodeAgent.runtime()
-            .increaseListenerReference(pipeMeta.getStaticMeta().getExtractorParameters());
-      }
     } catch (final Exception ex) {
       LOGGER.error("Failed to load pipe task info from snapshot", ex);
       loadPipeTaskInfoException = ex;

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/quota/QuotaInfo.java

@@ -258,6 +258,7 @@ public Map<String, TThrottleQuota> getThrottleQuotaLimit() {
 
   public void clear() {
     spaceQuotaLimit.clear();
+    spaceQuotaUsage.clear();
     throttleQuotaLimit.clear();
   }
 }

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/schema/TemplatePreSetTable.java

@@ -146,6 +146,8 @@ public void processLoadSnapshot(File snapshotDir) throws IOException {
     try {
       File snapshotFile = new File(snapshotDir, SNAPSHOT_FILENAME);
       if (!snapshotFile.exists()) {
+        // Empty preset tables are represented by the absence of a snapshot file.
+        templatePreSetMap.clear();
         return;
       }
 

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/schema/TemplateTable.java

@@ -260,16 +260,22 @@ public void processLoadSnapshot(File snapshotDir) throws IOException {
         BufferedInputStream bufferedInputStream = new BufferedInputStream(fileInputStream)) {
       // Load snapshot of template
       this.templateMap.clear();
+      this.templateIdMap.clear();
       deserialize(bufferedInputStream);
-      bufferedInputStream.close();
-      fileInputStream.close();
     } finally {
       templateReadWriteLock.writeLock().unlock();
     }
   }
 
   @TestOnly
   public void clear() {
-    this.templateMap.clear();
+    templateReadWriteLock.writeLock().lock();
+    try {
+      this.templateMap.clear();
+      this.templateIdMap.clear();
+      this.templateIdGenerator.set(0);
+    } finally {
+      templateReadWriteLock.writeLock().unlock();
+    }
   }
 }

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigRegionListenerSnapshotTest.java

@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.confignode.manager.pipe.agent.runtime;
+
+import org.apache.iotdb.commons.pipe.agent.task.meta.PipeMeta;
+import org.apache.iotdb.commons.pipe.agent.task.meta.PipeStaticMeta;
+import org.apache.iotdb.commons.pipe.config.constant.PipeSourceConstant;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+public class PipeConfigRegionListenerSnapshotTest {
+
+  @Test
+  public void testReconcileReferencesReplacesSnapshotState() throws Exception {
+    final Map<String, String> sourceAttributes = new HashMap<>();
+    sourceAttributes.put(PipeSourceConstant.EXTRACTOR_INCLUSION_KEY, "schema.database.create");
+    final PipeStaticMeta staticMeta =
+        new PipeStaticMeta(
+            "pipe", 1, sourceAttributes, Collections.emptyMap(), Collections.emptyMap());
+    final PipeMeta pipeMeta = Mockito.mock(PipeMeta.class);
+    Mockito.when(pipeMeta.getStaticMeta()).thenReturn(staticMeta);
+
+    final PipeConfigRegionListener listener = new PipeConfigRegionListener();
+    listener.increaseReference(staticMeta.getExtractorParameters());
+    listener.listener().close();
+
+    listener.reconcileReferences(Collections.emptyList());
+    listener.increaseReference(staticMeta.getExtractorParameters());
+    Assert.assertTrue(listener.listener().isOpened());
+
+    listener.reconcileReferences(Collections.singletonList(pipeMeta));
+    listener.reconcileReferences(Collections.singletonList(pipeMeta));
+    listener.decreaseReference(staticMeta.getExtractorParameters());
+    Assert.assertFalse(listener.listener().isOpened());
+  }
+}

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/QuotaInfoTest.java

@@ -37,6 +37,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -97,6 +98,19 @@ public void testSnapshot() throws TException, IOException {
     QuotaInfo quotaInfo2 = new QuotaInfo();
     quotaInfo2.processLoadSnapshot(snapshotDir);
 
+    final TSpaceQuota staleQuota = new TSpaceQuota();
+    staleQuota.setDeviceNum(1);
+    staleQuota.setTimeserieNum(1);
+    staleQuota.setDiskSize(1);
+    quotaInfo2.setSpaceQuota(
+        new SetSpaceQuotaPlan(Collections.singletonList("root.stale"), staleQuota));
+    Assert.assertTrue(quotaInfo2.getSpaceQuotaUsage().containsKey("root.stale"));
+
+    quotaInfo2.processLoadSnapshot(snapshotDir);
+
+    Assert.assertFalse(quotaInfo2.getSpaceQuotaLimit().containsKey("root.stale"));
+    Assert.assertFalse(quotaInfo2.getSpaceQuotaUsage().containsKey("root.stale"));
+
     Assert.assertEquals(quotaInfo.getSpaceQuotaLimit(), quotaInfo2.getSpaceQuotaLimit());
     Assert.assertEquals(quotaInfo.getThrottleQuotaLimit(), quotaInfo2.getThrottleQuotaLimit());
   }

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/schema/TemplatePreSetTableTest.java

@@ -104,11 +104,24 @@ public void testSnapshot() throws IllegalPathException {
       newTemplatePreSetTable = new TemplatePreSetTable();
       newTemplatePreSetTable.processLoadSnapshot(snapshotDir);
 
-      Assert.assertTrue(templatePreSetTable.isPreSet(templateId1, templateSetPath1));
-      Assert.assertTrue(templatePreSetTable.isPreSet(templateId2, templateSetPath1));
-      Assert.assertTrue(templatePreSetTable.isPreSet(templateId2, templateSetPath2));
+      Assert.assertTrue(newTemplatePreSetTable.isPreSet(templateId1, templateSetPath1));
+      Assert.assertTrue(newTemplatePreSetTable.isPreSet(templateId2, templateSetPath1));
+      Assert.assertTrue(newTemplatePreSetTable.isPreSet(templateId2, templateSetPath2));
     } catch (IOException e) {
       Assert.fail();
     }
   }
+
+  @Test
+  public void testEmptySnapshotReplacesExistingState() throws IOException, IllegalPathException {
+    final int templateId = 5;
+    final PartialPath templateSetPath = new PartialPath("root.db.t1");
+
+    Assert.assertTrue(templatePreSetTable.processTakeSnapshot(snapshotDir));
+    templatePreSetTable.preSetTemplate(templateId, templateSetPath);
+
+    templatePreSetTable.processLoadSnapshot(snapshotDir);
+
+    Assert.assertFalse(templatePreSetTable.isPreSet(templateId, templateSetPath));
+  }
 }

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/schema/TemplateTableTest.java

@@ -77,7 +77,8 @@ public void testSnapshot() throws IOException, MetadataException {
     }
 
     templateTable.processTakeSnapshot(snapshotDir);
-    templateTable.clear();
+    Template staleTemplate = newSchemaTemplate("stale_template");
+    templateTable.createTemplate(staleTemplate);
     templateTable.processLoadSnapshot(snapshotDir);
 
     // show nodes in schemaengine template
@@ -86,6 +87,13 @@ public void testSnapshot() throws IOException, MetadataException {
       Template template = templates.get(i);
       Assert.assertEquals(template, templateTable.getTemplate(templateNameTmp));
     }
+
+    try {
+      templateTable.getTemplate(staleTemplate.getId());
+      Assert.fail("Template created after the snapshot should be removed when loading it");
+    } catch (MetadataException expected) {
+      // expected
+    }
   }
 
   @Test

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSink.java

@@ -256,19 +256,24 @@ private void doTransfer(
     final List<File> sealedFiles = batchToTransfer.sealTsFiles();
     final Map<Pair<String, Long>, Double> pipe2WeightMap = batchToTransfer.deepCopyPipe2WeightMap();
 
-    for (final File tsFile : sealedFiles) {
-      doTransfer(pipe2WeightMap, socket, tsFile, null, null, tsFile.getName());
-      try {
-        RetryUtils.retryOnException(
-            () -> {
-              FileUtils.delete(tsFile);
-              return null;
-            });
-      } catch (final NoSuchFileException e) {
-        LOGGER.info("The file {} is not found, may already be deleted.", tsFile);
-      } catch (final Exception e) {
-        LOGGER.warn(
-            "Failed to delete batch file {}, this file should be deleted manually later", tsFile);
+    try {
+      for (final File tsFile : sealedFiles) {
+        doTransfer(pipe2WeightMap, socket, tsFile, null, null, tsFile.getName());
+      }
+    } finally {
+      for (final File tsFile : sealedFiles) {
+        try {
+          RetryUtils.retryOnException(
+              () -> {
+                FileUtils.delete(tsFile);
+                return null;
+              });
+        } catch (final NoSuchFileException e) {
+          LOGGER.info("The file {} is not found, may already be deleted.", tsFile);
+        } catch (final Exception e) {
+          LOGGER.warn(
+              "Failed to delete batch file {}, this file should be deleted manually later", tsFile);
+        }
       }
     }
   }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java

@@ -64,6 +64,7 @@
 import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq;
 
 import com.google.common.collect.ImmutableSet;
+import org.apache.commons.io.FileUtils;
 import org.apache.tsfile.exception.write.WriteProcessException;
 import org.apache.tsfile.utils.Pair;
 import org.slf4j.Logger;
@@ -249,6 +250,7 @@ private void transferInBatchWithoutCheck(
       final AtomicInteger eventsReferenceCount = new AtomicInteger(sealedFiles.size());
       final AtomicBoolean eventsHadBeenAddedToRetryQueue = new AtomicBoolean(false);
 
+      int transferredFileCount = 0;
       try {
         for (final File sealedFile : sealedFiles) {
           transfer(
@@ -262,8 +264,12 @@ private void transferInBatchWithoutCheck(
                   null,
                   false,
                   null));
+          transferredFileCount++;
         }
       } catch (final Exception e) {
+        for (int i = transferredFileCount; i < sealedFiles.size(); i++) {
+          FileUtils.deleteQuietly(sealedFiles.get(i));
+        }
         PipeLogger.log(LOGGER::warn, e, "Failed to transfer tsfile batch (%s).", sealedFiles);
         if (eventsHadBeenAddedToRetryQueue.compareAndSet(false, true)) {
           addFailureEventsToRetryQueue(events, e);
@@ -421,22 +427,28 @@ private boolean transferWithoutCheck(final TsFileInsertionEvent tsFileInsertionE
 
   private void transfer(final PipeTransferTsFileHandler pipeTransferTsFileHandler) {
     transferTsFileCounter.incrementAndGet();
-    CompletableFuture<Void> completableFuture =
-        CompletableFuture.supplyAsync(
-            () -> {
-              AsyncPipeDataTransferServiceClient client = null;
-              try {
-                client = transferTsFileClientManager.borrowClient();
-                pipeTransferTsFileHandler.transfer(transferTsFileClientManager, client);
-              } catch (final Exception ex) {
-                logOnClientException(client, ex);
-                pipeTransferTsFileHandler.onError(ex);
-              } finally {
-                transferTsFileCounter.decrementAndGet();
-              }
-              return null;
-            },
-            transferTsFileClientManager.getExecutor());
+    final CompletableFuture<Void> completableFuture;
+    try {
+      completableFuture =
+          CompletableFuture.supplyAsync(
+              () -> {
+                AsyncPipeDataTransferServiceClient client = null;
+                try {
+                  client = transferTsFileClientManager.borrowClient();
+                  pipeTransferTsFileHandler.transfer(transferTsFileClientManager, client);
+                } catch (final Exception ex) {
+                  logOnClientException(client, ex);
+                  pipeTransferTsFileHandler.onError(ex);
+                } finally {
+                  transferTsFileCounter.decrementAndGet();
+                }
+                return null;
+              },
+              transferTsFileClientManager.getExecutor());
+    } catch (final RuntimeException e) {
+      transferTsFileCounter.decrementAndGet();
+      throw e;
+    }
 
     if (PipeConfig.getInstance().isTransferTsFileSync()) {
       try {