[PIPE][Subscription] Fix pipe races and subscription WAL replay wakeup

integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/auto/enhanced/IoTDBPipeAutoDropIT.java

@@ -175,16 +175,8 @@ public void testAutoDropFinitePipesWithoutDataRegion() throws Exception {
                       .setExtractorAttributes(realtimeSourceAttributes))
               .getCode());
 
-      Assert.assertEquals(
-          TSStatusCode.SUCCESS_STATUS.getStatusCode(),
-          client.startPipe("query_pipe_without_data_region").getCode());
-      Assert.assertEquals(
-          TSStatusCode.SUCCESS_STATUS.getStatusCode(),
-          client.startPipe("history_pipe_without_data_region").getCode());
-      Assert.assertEquals(
-          TSStatusCode.SUCCESS_STATUS.getStatusCode(),
-          client.startPipe("realtime_pipe_without_data_region").getCode());
-
+      // Pipes start automatically after creation. Explicitly starting finite pipes can race with
+      // their auto-drop when they have no DataRegions to process.
       await()
           .pollInSameThread()
           .pollInterval(1L, TimeUnit.SECONDS)

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/PipeTsFileInsertionEvent.java

@@ -55,6 +55,7 @@
 import org.apache.iotdb.db.pipe.resource.tsfile.PipeTsFileResourceManager;
 import org.apache.iotdb.db.pipe.source.dataregion.realtime.assigner.PipeTsFileEpochProgressIndexKeeper;
 import org.apache.iotdb.db.storageengine.dataregion.memtable.TsFileProcessor;
+import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile;
 import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource;
 import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent;
 import org.apache.iotdb.pipe.api.event.dml.insertion.TsFileInsertionEvent;
@@ -274,7 +275,9 @@ private PipeTsFileInsertionEvent(
     // hard-link it to each pipe dir
     this.tsFile = Objects.isNull(tsFile) ? resource.getTsFile() : tsFile;
 
-    this.isWithMod = isWithMod && resource.anyModFileExists();
+    // Historical files and their mods have already been pinned. Compaction may have removed
+    // the original resource before this event is created, so inspect the pinned mods later.
+    this.isWithMod = isWithMod && (isGeneratedByHistoricalExtractor || resource.anyModFileExists());
     this.modFile = this.isWithMod ? resource.getExclusiveModFile().getFile() : null;
     // TODO: process the shared mod file
     this.sharedModFile =
@@ -463,9 +466,22 @@ public boolean internallyIncreaseResourceReferenceCount(final String holderMessa
         PipeTsFileResourceManager.getPipeTsFileResourcePipeName(pipeName, creationTime);
     final File originalTsFile = tsFile;
     final File originalModFile = modFile;
+    final boolean originalIsWithMod = isWithMod;
     File increasedTsFile = null;
     boolean increased = false;
     try {
+      if (isGeneratedByHistoricalExtractor) {
+        tsFile =
+            PipeTsFileResourceManager.getHardlinkOrCopiedFileInPipeDir(
+                tsFile, pipeTsFileResourcePipeName);
+        if (isWithMod) {
+          modFile = ModificationFile.getExclusiveMods(tsFile);
+          isWithMod = modFile.exists();
+          if (!isWithMod) {
+            modFile = null;
+          }
+        }
+      }
       increasedTsFile =
           PipeDataNodeResourceManager.tsfile()
               .increaseFileReference(tsFile, true, pipeTsFileResourcePipeName);
@@ -490,6 +506,7 @@ public boolean internallyIncreaseResourceReferenceCount(final String holderMessa
       }
       tsFile = originalTsFile;
       modFile = originalModFile;
+      isWithMod = originalIsWithMod;
       LOGGER.warn(
           String.format(
               DataNodePipeMessages.INCREASE_REFERENCE_COUNT_TSFILE_OR_MODFILE_ERROR_HOLDER_FMT,

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/tsfile/PipeTsFileResourceManager.java

@@ -427,14 +427,17 @@ public Map<String, TSDataType> getMeasurementDataTypeMapFromCache(
         : hardlinkOrCopiedFileToTsFilePublicResourceMap;
   }
 
-  public void pinTsFileResource(
+  /** Returns whether a companion mod-file reference was acquired. */
+  public boolean pinTsFileResource(
       final TsFileResource resource, final boolean withMods, final @Nullable String pipeName)
       throws IOException {
     final File pinnedTsFile = increaseFileReference(resource.getTsFile(), true, pipeName);
     try {
       if (withMods && resource.getExclusiveModFile().exists()) {
         increaseFileReference(resource.getExclusiveModFile().getFile(), false, pipeName);
+        return true;
       }
+      return false;
     } catch (final IOException | RuntimeException e) {
       // Pinning is a two-file operation.  Do not leave the TsFile pinned when the mod file cannot
       // be pinned (for example, when the pipe directory is temporarily unavailable).
@@ -448,9 +451,7 @@ public void pinTsFileResource(
   }
 
   public void unpinTsFileResource(
-      final TsFileResource resource,
-      final boolean shouldTransferModFile,
-      final @Nullable String pipeName)
+      final TsFileResource resource, final @Nullable String pipeName, final boolean pinnedModFile)
       throws IOException {
     Exception firstException = null;
     try {
@@ -462,7 +463,9 @@ public void unpinTsFileResource(
 
     // Always attempt the mod-file cleanup even when resolving/decreasing the TsFile fails.  A
     // failed first cleanup must not strand the second reference.
-    if (shouldTransferModFile && resource.exclusiveModFileExists()) {
+    // Use the state captured when pinning. Compaction can remove the original mod file, and
+    // an event can acquire mod-file references even when no mod file was originally pinned.
+    if (pinnedModFile) {
       try {
         decreaseFileReference(
             getHardlinkOrCopiedFileInPipeDir(resource.getExclusiveModFile().getFile(), pipeName),

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/source/dataregion/historical/PipeHistoricalDataRegionTsFileAndDeletionSource.java

@@ -175,6 +175,7 @@ public class PipeHistoricalDataRegionTsFileAndDeletionSource
   private Queue<PersistentResource> pendingQueue;
   private final Map<TsFileResource, Set<String>> filteredTsFileResources2TableNames =
       new HashMap<>();
+  private final Set<TsFileResource> resourcesWithPinnedModFile = new HashSet<>();
   private final Map<PersistentResource, Long> pendingResource2ReplicateIndexForIoTV2 =
       new HashMap<>();
   private final Set<PersistentResource> historicalProgressReportResources = new HashSet<>();
@@ -800,8 +801,11 @@ private void extractTsFiles(
                 // Pin the resource, in case the file is removed by compaction or anything.
                 // Will unpin it after the PipeTsFileInsertionEvent is created and pinned.
                 try {
-                  PipeDataNodeResourceManager.tsfile()
-                      .pinTsFileResource(resource, shouldTransferModFile, pipeNameWithCreationTime);
+                  if (PipeDataNodeResourceManager.tsfile()
+                      .pinTsFileResource(
+                          resource, shouldTransferModFile, pipeNameWithCreationTime)) {
+                    resourcesWithPinnedModFile.add(resource);
+                  }
                   return false;
                 } catch (final IOException e) {
                   ++statistics.pinFailedCount;
@@ -1217,7 +1221,8 @@ protected boolean consumeSkippedHistoricalTsFileEventIfNecessary(final TsFileRes
     } finally {
       try {
         PipeDataNodeResourceManager.tsfile()
-            .unpinTsFileResource(resource, shouldTransferModFile, pipeNameWithCreationTime);
+            .unpinTsFileResource(
+                resource, pipeNameWithCreationTime, resourcesWithPinnedModFile.remove(resource));
       } catch (final IOException e) {
         LOGGER.warn(
             DataNodePipeMessages.PIPE_FAILED_TO_UNPIN_SKIPPED_HISTORICAL_TSFILERESOURCE,
@@ -1329,7 +1334,8 @@ protected Event supplyTsFileEvent(final TsFileResource resource) {
       if (shouldUnpinResource) {
         try {
           PipeDataNodeResourceManager.tsfile()
-              .unpinTsFileResource(resource, shouldTransferModFile, pipeNameWithCreationTime);
+              .unpinTsFileResource(
+                  resource, pipeNameWithCreationTime, resourcesWithPinnedModFile.remove(resource));
         } catch (final IOException e) {
           LOGGER.warn(
               DataNodePipeMessages.PIPE_FAILED_TO_UNPIN_TSFILERESOURCE_AFTER_CREATING,
@@ -1454,7 +1460,9 @@ public synchronized void close() {
               try {
                 PipeDataNodeResourceManager.tsfile()
                     .unpinTsFileResource(
-                        (TsFileResource) resource, shouldTransferModFile, pipeNameWithCreationTime);
+                        (TsFileResource) resource,
+                        pipeNameWithCreationTime,
+                        resourcesWithPinnedModFile.remove(resource));
               } catch (final IOException e) {
                 LOGGER.warn(
                     DataNodePipeMessages.PIPE_FAILED_TO_UNPIN_TSFILERESOURCE_AFTER_DROPPING,
@@ -1468,6 +1476,7 @@ public synchronized void close() {
       pendingQueue = null;
     }
     pendingResource2ReplicateIndexForIoTV2.clear();
+    resourcesWithPinnedModFile.clear();
     historicalProgressReportResources.clear();
     pendingHistoricalProgressIndexToReport = null;
   }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TableConfigTaskVisitor.java

@@ -1370,6 +1370,15 @@ public static void checkAndEnrichSourceUser(
                   .FAILED_TO_S_PIPE_S_IN_IOTDB_SOURCE_PASSWORD_MUST_BE_SET_WHEN_THE_USERNAME_IS_SPECIFIED,
               isAlter ? DataNodeQueryMessages.ALTER : DataNodeQueryMessages.CREATE,
               pipeName));
+    } else {
+      updatePipeUserIdFromUsername(
+          replacedSourceAttributes,
+          sourceParameters,
+          PipeSourceConstant.SOURCE_IOTDB_USER_ID,
+          PipeSourceConstant.EXTRACTOR_IOTDB_USER_KEY,
+          PipeSourceConstant.SOURCE_IOTDB_USER_KEY,
+          PipeSourceConstant.EXTRACTOR_IOTDB_USERNAME_KEY,
+          PipeSourceConstant.SOURCE_IOTDB_USERNAME_KEY);
     }
   }
 
@@ -1442,6 +1451,15 @@ public static void checkAndEnrichSinkUser(
                   .FAILED_TO_S_PIPE_S_IN_WRITE_BACK_SINK_PASSWORD_MUST_BE_SET_WHEN_THE_USERNAME_IS,
               isAlter ? DataNodeQueryMessages.ALTER : DataNodeQueryMessages.CREATE,
               pipeName));
+    } else {
+      updatePipeUserIdFromUsername(
+          connectorAttributes,
+          connectorParameters,
+          PipeSinkConstant.SINK_IOTDB_USER_ID,
+          PipeSinkConstant.CONNECTOR_IOTDB_USER_KEY,
+          PipeSinkConstant.SINK_IOTDB_USER_KEY,
+          PipeSinkConstant.CONNECTOR_IOTDB_USERNAME_KEY,
+          PipeSinkConstant.SINK_IOTDB_USERNAME_KEY);
     }
   }
 
@@ -1510,6 +1528,14 @@ public static void markSourceAuthenticationAsExplicitIfNecessary(
         PipeSourceConstant.SOURCE_IOTDB_PASSWORD_KEY)) {
       sourceAttributes.put(
           SystemConstant.SOURCE_AUTHENTICATION_INJECTED_KEY, Boolean.FALSE.toString());
+      updatePipeUserIdFromUsername(
+          sourceAttributes,
+          sourceParameters,
+          PipeSourceConstant.SOURCE_IOTDB_USER_ID,
+          PipeSourceConstant.EXTRACTOR_IOTDB_USER_KEY,
+          PipeSourceConstant.SOURCE_IOTDB_USER_KEY,
+          PipeSourceConstant.EXTRACTOR_IOTDB_USERNAME_KEY,
+          PipeSourceConstant.SOURCE_IOTDB_USERNAME_KEY);
     }
   }
 
@@ -1524,6 +1550,26 @@ public static void markSinkAuthenticationAsExplicitIfNecessary(
         PipeSinkConstant.CONNECTOR_IOTDB_PASSWORD_KEY,
         PipeSinkConstant.SINK_IOTDB_PASSWORD_KEY)) {
       sinkAttributes.put(SystemConstant.SINK_AUTHENTICATION_INJECTED_KEY, Boolean.FALSE.toString());
+      updatePipeUserIdFromUsername(
+          sinkAttributes,
+          sinkParameters,
+          PipeSinkConstant.SINK_IOTDB_USER_ID,
+          PipeSinkConstant.CONNECTOR_IOTDB_USER_KEY,
+          PipeSinkConstant.SINK_IOTDB_USER_KEY,
+          PipeSinkConstant.CONNECTOR_IOTDB_USERNAME_KEY,
+          PipeSinkConstant.SINK_IOTDB_USERNAME_KEY);
+    }
+  }
+
+  private static void updatePipeUserIdFromUsername(
+      final Map<String, String> attributes,
+      final PipeParameters parameters,
+      final String userIdAttributeKey,
+      final String... usernameAttributeKeys) {
+    final String username = parameters.getStringByKeys(usernameAttributeKeys);
+    if (username != null) {
+      attributes.put(
+          userIdAttributeKey, String.valueOf(AuthorityChecker.getUserId(username).orElse(-1L)));
     }
   }
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java

@@ -518,34 +518,39 @@ private boolean reject(final SubscriptionQueueRejectionReason reason) {
           Objects.isNull(reason) ? SubscriptionQueueRejectionReason.QUEUE_CAPACITY : reason;
       lastOfferRejectionReason.set(effectiveReason);
       rejectionListener.accept(effectiveReason);
-      return false;
+      return effectiveReason != SubscriptionQueueRejectionReason.INACTIVE_OR_CLOSED;
     }
 
     @Override
     public boolean offer(final IndexedConsensusRequest request) {
       final boolean offered;
+      boolean wakeupRequired = false;
       synchronized (this) {
         final SubscriptionQueueRejectionReason admissionReason = admissionSupplier.get();
         if (admissionReason != SubscriptionQueueRejectionReason.NONE) {
-          return reject(admissionReason);
-        }
-        if (!requestMemoryManager.reserve(request)) {
-          return reject(SubscriptionQueueRejectionReason.CONSENSUS_REQUEST_MEMORY_LIMIT);
-        }
-        try {
-          offered = super.offer(request);
-        } catch (final Throwable t) {
-          requestMemoryManager.free(request);
-          throw t;
-        }
-        if (offered) {
-          retainedRequestBytes.addAndGet(request.getRetainedMemorySize());
+          wakeupRequired = reject(admissionReason);
+          offered = false;
+        } else if (!requestMemoryManager.reserve(request)) {
+          wakeupRequired = reject(SubscriptionQueueRejectionReason.CONSENSUS_REQUEST_MEMORY_LIMIT);
+          offered = false;
         } else {
-          requestMemoryManager.free(request);
-          reject(SubscriptionQueueRejectionReason.QUEUE_CAPACITY);
+          try {
+            offered = super.offer(request);
+          } catch (final Throwable t) {
+            requestMemoryManager.free(request);
+            throw t;
+          }
+          if (offered) {
+            retainedRequestBytes.addAndGet(request.getRetainedMemorySize());
+          } else {
+            requestMemoryManager.free(request);
+            wakeupRequired = reject(SubscriptionQueueRejectionReason.QUEUE_CAPACITY);
+          }
         }
       }
-      if (offered) {
+      // The rejected request remains available in WAL. Wake the replay worker so a queue that has
+      // already gone dormant does not wait for a later successful realtime offer to make progress.
+      if (offered || wakeupRequired) {
         wakeupHook.run();
       }
       return offered;