Hold the consumer semaphore while a replace reverts upsert metadata

pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java

@@ -1207,7 +1207,9 @@ SegmentBuildDescriptor getSegmentBuildDescriptor() {
   @Nullable
   protected SegmentBuildDescriptor buildSegmentInternal(boolean forCommit)
       throws SegmentBuildFailureException {
-    if (_parallelSegmentConsumptionPolicy.isAllowedDuringBuild()) {
+    // Not released when the replace rewrites upsert or dedup metadata: a build that fails or mismatches the
+    // committed CRC falls back to a download, which cannot hold back a semaphore already given away.
+    if (_parallelSegmentConsumptionPolicy.isAllowedDuringBuild() && !holdsConsumerSemaphoreForMetadata()) {
       closeStreamConsumerAndReleaseSemaphore();
     }
     // Do not allow building segment when table data manager is already shut down
@@ -1765,26 +1767,40 @@ public void goOnlineFromConsuming(SegmentZKMetadata segmentZKMetadata)
 
   protected void downloadSegmentAndReplace(SegmentZKMetadata segmentZKMetadata)
       throws Exception {
-    boolean allowedDuringDownload = _parallelSegmentConsumptionPolicy.isAllowedDuringDownload();
     // A local build that failed or mismatched the committed CRC has already released the semaphore when the policy
     // allows consumption during build, so the next consuming segment can be running during this download as well
     boolean releasedDuringBuild =
         _parallelSegmentConsumptionPolicy.isAllowedDuringBuild() && !_consumerSemaphoreAcquired.get();
-    if ((allowedDuringDownload || releasedDuringBuild) && _partitionUpsertMetadataManager != null
-        && _partitionUpsertMetadataManager.shouldRevertMetadataOnInconsistency(_realtimeSegment)) {
-      // Table config validation rejects this in PROTECTED mode. It still happens when the mode is switched after the
-      // table is created, when the server-level default allows partial upsert consumption during commit, or when a
-      // pauseless table falls back to a download after its local build fails.
-      _segmentLogger.error("Next consuming segment can run during download while this replace reverts upsert "
-          + "metadata. Its snapshot can miss the rows the revert restores. Use DISALLOW_ALWAYS");
-      _serverMetrics.addMeteredTableValue(_clientId, ServerMeter.UPSERT_REVERT_WITH_CONSUMPTION_DURING_DOWNLOAD, 1L);
+    // Hold the semaphore when the replace rewrites the key map, so the next consuming segment does not start against
+    // a partially reverted one, and does not replay rows against keys this replace is about to remove.
+    boolean holdsConsumerSemaphore = holdsConsumerSemaphoreForMetadata();
+    boolean allowedDuringDownload =
+        _parallelSegmentConsumptionPolicy.isAllowedDuringDownload() && !holdsConsumerSemaphore;
+    if (releasedDuringBuild && holdsConsumerSemaphore) {
+      // The build holds the semaphore for these tables, so reaching this means it was released on another path.
+      _segmentLogger.error("Next consuming segment can run during download while this replace rewrites upsert or "
+          + "dedup metadata. Its rows can be dropped against keys this replace removes. Use DISALLOW_ALWAYS");
+      if (revertsUpsertMetadata()) {
+        _serverMetrics.addMeteredTableValue(_clientId, ServerMeter.UPSERT_REVERT_WITH_CONSUMPTION_DURING_DOWNLOAD, 1L);
+      }
     }
     if (allowedDuringDownload) {
       closeStreamConsumerAndReleaseSemaphore();
     }
     _realtimeTableDataManager.downloadAndReplaceConsumingSegment(segmentZKMetadata);
   }
 
+  /// The replace rewrites the key map, so hold the semaphore across the build and the download like
+  /// DISALLOW_ALWAYS: upsert reverts those keys, dedup removes them when the mutable segment is offloaded.
+  private boolean holdsConsumerSemaphoreForMetadata() {
+    return revertsUpsertMetadata() || _partitionDedupMetadataManager != null;
+  }
+
+  private boolean revertsUpsertMetadata() {
+    return _partitionUpsertMetadataManager != null
+        && _partitionUpsertMetadataManager.shouldRevertMetadataOnInconsistency(_realtimeSegment);
+  }
+
   protected long now() {
     return System.currentTimeMillis();
   }

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/UpsertContext.java

@@ -26,6 +26,7 @@
 import javax.annotation.Nullable;
 import org.apache.commons.collections4.CollectionUtils;
 import org.apache.commons.lang3.builder.ToStringBuilder;
+import org.apache.pinot.common.utils.PauselessConsumptionUtils;
 import org.apache.pinot.segment.local.data.manager.TableDataManager;
 import org.apache.pinot.spi.config.table.HashFunction;
 import org.apache.pinot.spi.config.table.TableConfig;
@@ -201,8 +202,11 @@ public File getTableIndexDir() {
   /// during segment replacement after force commit. This happens when:
   /// - Partial upsert is enabled (records need to be merged with previous values)
   /// - dropOutOfOrderRecord is enabled with NONE consistency mode (records may have been dropped)
+  /// Pauseless tables are excluded: they consume while the previous segment is still building, so the upsert
+  /// metadata cannot be kept reconcilable across replicas.
   public boolean isTableTypeInconsistentDuringConsumption() {
-    return _dropOutOfOrderRecord || _outOfOrderRecordColumn != null || _partialUpsertHandlerSupplier != null;
+    return (_dropOutOfOrderRecord || _outOfOrderRecordColumn != null || _partialUpsertHandlerSupplier != null)
+        && !PauselessConsumptionUtils.isPauselessEnabled(_tableConfig);
   }
 
   @Override