Add an opt-in controller-copy METADATA push for same-name segment refreshes

pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutor.java

@@ -26,7 +26,9 @@
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import javax.annotation.Nullable;
 import org.apache.commons.io.FileUtils;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.hc.core5.http.Header;
 import org.apache.hc.core5.http.HttpHeaders;
 import org.apache.hc.core5.http.NameValuePair;
@@ -142,10 +144,27 @@ public SegmentConversionResult executeTask(PinotTaskConfig pinotTaskConfig)
         reportTaskProcessingMetrics(tableNameWithType, taskType, segmentMetadata.getTotalDocs());
       }
 
+      BatchConfigProperties.SegmentPushType pushType = getSegmentPushType(configs);
+      boolean copyToDeepStore =
+          pushType == BatchConfigProperties.SegmentPushType.METADATA && isCopyToDeepStoreForMetadataPush();
+      // Controller-copy pushes send only metadata.properties and creation.meta, built from the local converted
+      // segment. An unchanged segment (same CRC) skips the tar and the staging and only re-registers its metadata.
+      File segmentMetadataTarFile = null;
+      boolean reuseExistingSegment = false;
+      if (copyToDeepStore) {
+        segmentMetadataTarFile = createSegmentMetadataTarFile(convertedSegmentDir, tempDataDir, segmentName);
+        long convertedSegmentCrc = Long.parseLong(new SegmentMetadataImpl(convertedSegmentDir).getCrc());
+        reuseExistingSegment =
+            convertedSegmentCrc == Long.parseLong(originalSegmentCrc) && StringUtils.isNotEmpty(downloadURL);
+      }
+
       // Tar the converted segment
-      _eventObserver.notifyProgress(_pinotTaskConfig, "Compressing segment: " + segmentName);
-      File convertedTarredSegmentFile = new File(tempDataDir, segmentName + TarCompressionUtils.TAR_GZ_FILE_EXTENSION);
-      TarCompressionUtils.createCompressedTarFile(convertedSegmentDir, convertedTarredSegmentFile);
+      File convertedTarredSegmentFile = null;
+      if (!reuseExistingSegment) {
+        _eventObserver.notifyProgress(_pinotTaskConfig, "Compressing segment: " + segmentName);
+        convertedTarredSegmentFile = new File(tempDataDir, segmentName + TarCompressionUtils.TAR_GZ_FILE_EXTENSION);
+        TarCompressionUtils.createCompressedTarFile(convertedSegmentDir, convertedTarredSegmentFile);
+      }
       if (!FileUtils.deleteQuietly(convertedSegmentDir)) {
         LOGGER.warn("Failed to delete converted segment: {}", convertedSegmentDir.getAbsolutePath());
       }
@@ -180,6 +199,7 @@ public SegmentConversionResult executeTask(PinotTaskConfig pinotTaskConfig)
           new BasicHeader(FileUploadDownloadClient.CustomHeaders.SEGMENT_ZK_METADATA_CUSTOM_MAP_MODIFIER,
               segmentZKMetadataCustomMapModifier.toJsonString());
 
+      // Both push modes send the same guards: IF-MATCH (segment refreshed since) and REFRESH_ONLY (segment deleted).
       List<Header> httpHeaders = new ArrayList<>();
       httpHeaders.add(ifMatchHeader);
       httpHeaders.add(refreshOnlyHeader);
@@ -189,8 +209,7 @@ public SegmentConversionResult executeTask(PinotTaskConfig pinotTaskConfig)
       // Set parameters for upload request (shared with metadata push).
       List<NameValuePair> parameters = getSegmentPushCommonParams(tableNameWithType);
 
-      // Upload the tarred segment using the configured push mode (TAR or METADATA)
-      BatchConfigProperties.SegmentPushType pushType = getSegmentPushType(configs);
+      // Upload the segment using the configured push mode (TAR or METADATA)
       _eventObserver.notifyProgress(_pinotTaskConfig, "Uploading segment: " + segmentName + " (push mode: " + pushType
           + ")");
       try {
@@ -200,8 +219,13 @@ public SegmentConversionResult executeTask(PinotTaskConfig pinotTaskConfig)
                 uploadURL, convertedTarredSegmentFile);
             break;
           case METADATA:
-            uploadSegmentWithMetadata(configs, pinotTaskConfig, segmentConversionResult, authProvider, parameters,
-                tableNameWithType, convertedTarredSegmentFile);
+            if (copyToDeepStore) {
+              uploadSegmentMetadataWithControllerCopy(configs, httpHeaders, parameters, tableNameWithType, segmentName,
+                  uploadURL, downloadURL, convertedTarredSegmentFile, segmentMetadataTarFile);
+            } else {
+              uploadSegmentWithMetadata(configs, pinotTaskConfig, segmentConversionResult, authProvider, parameters,
+                  tableNameWithType, convertedTarredSegmentFile);
+            }
             break;
           default:
             throw new UnsupportedOperationException("Unrecognized push mode: " + pushType);
@@ -212,9 +236,12 @@ public SegmentConversionResult executeTask(PinotTaskConfig pinotTaskConfig)
         _eventObserver.notifyTaskError(_pinotTaskConfig, e);
         throw e;
       } finally {
-        if (!FileUtils.deleteQuietly(convertedTarredSegmentFile)) {
+        if (convertedTarredSegmentFile != null && !FileUtils.deleteQuietly(convertedTarredSegmentFile)) {
           LOGGER.warn("Failed to delete tarred converted segment: {}", convertedTarredSegmentFile.getAbsolutePath());
         }
+        if (segmentMetadataTarFile != null) {
+          FileUtils.deleteQuietly(segmentMetadataTarFile);
+        }
       }
 
       LOGGER.info("Done executing {} on table: {}, segment: {}", taskType, tableNameWithType, segmentName);
@@ -224,6 +251,12 @@ public SegmentConversionResult executeTask(PinotTaskConfig pinotTaskConfig)
     }
   }
 
+  /// Opt-in for tasks that refresh a segment under its own name: stage under a task-unique name, keep the TAR guards
+  /// and let the controller copy the tar into the segment's deep store location. Default false keeps current behavior.
+  protected boolean isCopyToDeepStoreForMetadataPush() {
+    return false;
+  }
+
   /// Pushes the segment in METADATA (or URI) mode: copies the tarred segment to the output PinotFS and sends segment
   /// URI and metadata to the controller. Requires [BatchConfigProperties#OUTPUT_SEGMENT_DIR_URI] and
   /// [BatchConfigProperties#PUSH_CONTROLLER_URI] in configs.
@@ -268,6 +301,61 @@ private void uploadSegmentWithMetadata(Map<String, String> configs, PinotTaskCon
     }
   }
 
+  /// METADATA push for a same-name refresh: stage at `<segment>.<taskId>.tar.gz` (never a live path), let the
+  /// controller copy it into place, then always delete the staged tar. A null tar means the segment is unchanged.
+  private void uploadSegmentMetadataWithControllerCopy(Map<String, String> configs, List<Header> httpHeaders,
+      List<NameValuePair> parameters, String tableNameWithType, String segmentName, String uploadURL,
+      String downloadURL, @Nullable File convertedTarredSegmentFile, File segmentMetadataTarFile)
+      throws Exception {
+    if (convertedTarredSegmentFile == null) {
+      LOGGER.info("Segment: {} of table: {} is unchanged, registering metadata against: {}", segmentName,
+          tableNameWithType, downloadURL);
+      SegmentConversionUtils.uploadSegmentMetadata(configs, withMetadataPushHeaders(httpHeaders, downloadURL, false),
+          parameters, tableNameWithType, segmentName, uploadURL, segmentMetadataTarFile);
+      return;
+    }
+    if (!configs.containsKey(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI)) {
+      throw new RuntimeException("Output dir URI missing for metadata push. Set "
+          + BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI + " in task config.");
+    }
+    URI stagedSegmentTarURI = moveSegmentToOutputPinotFS(configs, convertedTarredSegmentFile,
+        getStagedSegmentTarName(segmentName), true);
+    LOGGER.info("Staged converted segment: {} of table: {} at: {}", segmentName, tableNameWithType,
+        stagedSegmentTarURI);
+    try {
+      SegmentConversionUtils.uploadSegmentMetadata(configs,
+          withMetadataPushHeaders(httpHeaders, toSchemeQualifiedUri(stagedSegmentTarURI), true), parameters,
+          tableNameWithType, segmentName, uploadURL, segmentMetadataTarFile);
+    } finally {
+      deleteFromOutputPinotFS(configs, stagedSegmentTarURI);
+    }
+  }
+
+  /// The controller picks the PinotFS by scheme, so a plain-path (local) output dir is sent as a file URI.
+  private static String toSchemeQualifiedUri(URI uri) {
+    return uri.getScheme() == null ? new File(uri.getPath()).toURI().toString() : uri.toString();
+  }
+
+  /// Task-unique staging name: no collision across tasks, and a retry of the same task overwrites its leftover. It
+  /// stays in the table dir, so a leftover from a dead minion is an untracked segment for RetentionManager's sweep.
+  @VisibleForTesting
+  String getStagedSegmentTarName(String segmentName) {
+    String taskId = _pinotTaskConfig.getTaskId();
+    return segmentName + "." + (taskId != null ? taskId : UUID.randomUUID())
+        + TarCompressionUtils.TAR_GZ_FILE_EXTENSION;
+  }
+
+  private static List<Header> withMetadataPushHeaders(List<Header> httpHeaders, String segmentTarURI,
+      boolean copyToDeepStore) {
+    List<Header> headers = new ArrayList<>(httpHeaders);
+    headers.add(new BasicHeader(FileUploadDownloadClient.CustomHeaders.DOWNLOAD_URI, segmentTarURI));
+    headers.add(new BasicHeader(FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE,
+        FileUploadDownloadClient.FileUploadType.METADATA.toString()));
+    headers.add(new BasicHeader(FileUploadDownloadClient.CustomHeaders.COPY_SEGMENT_TO_DEEP_STORE,
+        String.valueOf(copyToDeepStore)));
+    return headers;
+  }
+
   // For tests only.
   @VisibleForTesting
   public void setMinionEventObserver(MinionEventObserver observer) {

pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseTaskExecutor.java

@@ -20,6 +20,7 @@
 
 import com.google.common.base.Preconditions;
 import java.io.File;
+import java.io.IOException;
 import java.net.URI;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -44,12 +45,14 @@
 import org.apache.pinot.core.util.PeerServerSegmentFinder;
 import org.apache.pinot.minion.MinionContext;
 import org.apache.pinot.minion.executor.PinotTaskExecutor;
+import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths;
 import org.apache.pinot.spi.auth.AuthProvider;
 import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.config.table.TableType;
 import org.apache.pinot.spi.data.Schema;
 import org.apache.pinot.spi.filesystem.PinotFS;
 import org.apache.pinot.spi.ingestion.batch.BatchConfigProperties;
+import org.apache.pinot.spi.ingestion.batch.spec.Constants;
 import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
 import org.apache.pinot.spi.ingestion.batch.spec.PushJobSpec;
 import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
@@ -210,18 +213,25 @@ protected SegmentGenerationJobSpec generateSegmentGenerationJobSpec(String table
     return spec;
   }
 
-  /// Copies the local segment tar file to the output PinotFS. Requires
-  /// [BatchConfigProperties#OUTPUT_SEGMENT_DIR_URI] in configs.
+  /// Copies the local segment tar to the output PinotFS under its own name. Fails if the target exists unless
+  /// [BatchConfigProperties#OVERWRITE_OUTPUT] is set.
   ///
   /// @return the URI of the segment tar on the output filesystem
   protected URI moveSegmentToOutputPinotFS(Map<String, String> configs, File localSegmentTarFile)
       throws Exception {
+    return moveSegmentToOutputPinotFS(configs, localSegmentTarFile, localSegmentTarFile.getName(),
+        Boolean.parseBoolean(configs.get(BatchConfigProperties.OVERWRITE_OUTPUT)));
+  }
+
+  /// Copies the local segment tar to `<outputDir>/<outputFileName>`, replacing an existing file only if asked.
+  protected URI moveSegmentToOutputPinotFS(Map<String, String> configs, File localSegmentTarFile,
+      String outputFileName, boolean overwrite)
+      throws Exception {
     URI outputSegmentDirURI = URI.create(configs.get(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI));
     try (PinotFS outputFileFS = MinionTaskUtils.getOutputPinotFS(configs, outputSegmentDirURI)) {
       URI outputSegmentTarURI = URI.create(MinionTaskUtils.normalizeDirectoryURI(outputSegmentDirURI)
-          + URIUtils.encode(localSegmentTarFile.getName()));
-      if (!Boolean.parseBoolean(configs.get(BatchConfigProperties.OVERWRITE_OUTPUT))
-          && outputFileFS.exists(outputSegmentTarURI)) {
+          + URIUtils.encode(outputFileName));
+      if (!overwrite && outputFileFS.exists(outputSegmentTarURI)) {
         throw new RuntimeException("Output file: " + outputSegmentTarURI + " already exists. Set 'overwriteOutput' to "
             + "true to ignore this error");
       }
@@ -230,6 +240,31 @@ protected URI moveSegmentToOutputPinotFS(Map<String, String> configs, File local
     }
   }
 
+  /// Best-effort delete on the output PinotFS. Failures are logged, so cleanup never masks the push outcome.
+  protected void deleteFromOutputPinotFS(Map<String, String> configs, URI fileURI) {
+    try (PinotFS outputFileFS = MinionTaskUtils.getOutputPinotFS(configs, fileURI)) {
+      outputFileFS.delete(fileURI, true);
+    } catch (Exception e) {
+      LOGGER.warn("Failed to delete: {} from the output PinotFS", fileURI, e);
+    }
+  }
+
+  /// Tars only metadata.properties and creation.meta from the local segment, which is all a METADATA push sends.
+  protected File createSegmentMetadataTarFile(File segmentDir, File outputDir, String segmentName)
+      throws IOException {
+    File metadataDir = new File(outputDir, segmentName + "-metadata");
+    File metadataTarFile = new File(outputDir, segmentName + Constants.METADATA_TAR_GZ_FILE_EXT);
+    try {
+      FileUtils.forceMkdir(metadataDir);
+      FileUtils.copyFileToDirectory(SegmentDirectoryPaths.findMetadataFile(segmentDir), metadataDir);
+      FileUtils.copyFileToDirectory(SegmentDirectoryPaths.findCreationMetaFile(segmentDir), metadataDir);
+      TarCompressionUtils.createCompressedTarFile(metadataDir, metadataTarFile);
+      return metadataTarFile;
+    } finally {
+      FileUtils.deleteQuietly(metadataDir);
+    }
+  }
+
   /// Returns HTTP parameters common to segment upload and metadata push (parallel push protection, table name, type).
   protected List<NameValuePair> getSegmentPushCommonParams(String tableNameWithType) {
     List<NameValuePair> params = new ArrayList<>();

pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/SegmentConversionUtils.java

@@ -108,6 +108,31 @@ public static Set<String> getSegmentNamesForTable(String tableNameWithType, URI
   public static void uploadSegment(Map<String, String> configs, List<Header> httpHeaders,
       List<NameValuePair> parameters, String tableNameWithType, String segmentName, String uploadURL, File fileToUpload)
       throws Exception {
+    uploadWithRetry(configs, httpHeaders, tableNameWithType, segmentName, uploadURL,
+        (client, uri, socketTimeoutMs) -> client.uploadSegment(uri, segmentName, fileToUpload, httpHeaders, parameters,
+            socketTimeoutMs));
+  }
+
+  /// METADATA-mode registration with the same retry loop as [#uploadSegment]. Only the metadata tar is sent,
+  /// and `httpHeaders` must carry the segment tar location as DOWNLOAD_URI.
+  public static void uploadSegmentMetadata(Map<String, String> configs, List<Header> httpHeaders,
+      List<NameValuePair> parameters, String tableNameWithType, String segmentName, String uploadURL,
+      File segmentMetadataFile)
+      throws Exception {
+    uploadWithRetry(configs, httpHeaders, tableNameWithType, segmentName, uploadURL,
+        (client, uri, socketTimeoutMs) -> client.uploadSegmentMetadata(uri, segmentName, segmentMetadataFile,
+            httpHeaders, parameters, socketTimeoutMs));
+  }
+
+  @FunctionalInterface
+  private interface UploadRequest {
+    SimpleHttpResponse send(FileUploadDownloadClient client, URI uri, int socketTimeoutMs)
+        throws Exception;
+  }
+
+  private static void uploadWithRetry(Map<String, String> configs, List<Header> httpHeaders, String tableNameWithType,
+      String segmentName, String uploadURL, UploadRequest uploadRequest)
+      throws Exception {
     // Create a RoundRobinURIProvider to round-robin IP addresses when retry uploading. Otherwise, it may always try to
     // upload to a same broken host as: 1) DNS may not RR the IP addresses 2) OS cache the DNS resolution result.
     RoundRobinURIProvider uriProvider = new RoundRobinURIProvider(List.of(new URI(uploadURL)), true);
@@ -145,9 +170,7 @@ public static void uploadSegment(Map<String, String> configs, List<Header> httpH
           httpHeaders.add(new BasicHeader(HttpHeaders.HOST, hostName + ":" + hostPort));
         }
         try {
-          SimpleHttpResponse response =
-              fileUploadDownloadClient.uploadSegment(uri, segmentName, fileToUpload, httpHeaders, parameters,
-                  socketTimeoutMs);
+          SimpleHttpResponse response = uploadRequest.send(fileUploadDownloadClient, uri, socketTimeoutMs);
           LOGGER.info("Got response {}: {} while uploading table: {}, segment: {} with uploadURL: {}",
               response.getStatusCode(), response.getResponse(), tableNameWithType, segmentName, uploadURL);
           return true;