From c0a33a75158e40952ee9d37a67827c0520a3c543 Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Mon, 7 Sep 2026 14:24:38 +0530 Subject: [PATCH] Add an opt-in controller-copy METADATA push for same-name segment refreshes The METADATA mode of BaseSingleSegmentConversionExecutor registers the staged tar's URI as the segment's download URL. For a task that refreshes a segment under its own name that path collides with a live segment stored as .tar.gz, drops the TAR path's If-Match and REFRESH_ONLY guards, and moves the download URL away from the deep store location so the second refresh of the same segment fails with "already exists". Add a protected hook, isCopyToDeepStoreForMetadataPush(), default false, so the existing behavior is unchanged for every current task. A task that returns true gets a refresh-safe path: the tar is staged under ..tar.gz, the push carries the same If-Match, REFRESH_ONLY and custom-map headers as TAR plus COPY_SEGMENT_TO_DEEP_STORE=true, the controller copies the bytes into the segment's existing deep store location, and the staged tar is deleted in a finally. The metadata tar is built from the local converted segment instead of downloading the staged tar back, and an unchanged segment (same CRC) only re-registers its metadata against the current download URL. SegmentConversionUtils gains uploadSegmentMetadata sharing the retry loop with uploadSegment. BaseTaskExecutor gains a moveSegmentToOutputPinotFS overload with an explicit target name and overwrite flag, deleteFromOutputPinotFS and createSegmentMetadataTarFile. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Y375AgHYsh1YqNSsfvF1a8 --- .../BaseSingleSegmentConversionExecutor.java | 104 ++++++- .../plugin/minion/tasks/BaseTaskExecutor.java | 45 ++- .../minion/tasks/SegmentConversionUtils.java | 29 +- ...seSingleSegmentConversionExecutorTest.java | 281 ++++++++++++++++-- 4 files changed, 415 insertions(+), 44 deletions(-) diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutor.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutor.java index 51ad6d2ef936..2da32aa7ff71 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutor.java +++ b/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
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 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 configs, PinotTaskCon } } + /// METADATA push for a same-name refresh: stage at `..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 configs, List
httpHeaders, + List 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
withMetadataPushHeaders(List
httpHeaders, String segmentTarURI, + boolean copyToDeepStore) { + List
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) { diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseTaskExecutor.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseTaskExecutor.java index 03eee9835767..61fe7b746015 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseTaskExecutor.java +++ b/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 configs, File localSegmentTarFile) throws Exception { + return moveSegmentToOutputPinotFS(configs, localSegmentTarFile, localSegmentTarFile.getName(), + Boolean.parseBoolean(configs.get(BatchConfigProperties.OVERWRITE_OUTPUT))); + } + + /// Copies the local segment tar to `/`, replacing an existing file only if asked. + protected URI moveSegmentToOutputPinotFS(Map 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 configs, File local } } + /// Best-effort delete on the output PinotFS. Failures are logged, so cleanup never masks the push outcome. + protected void deleteFromOutputPinotFS(Map 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 getSegmentPushCommonParams(String tableNameWithType) { List params = new ArrayList<>(); diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/SegmentConversionUtils.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/SegmentConversionUtils.java index 7d023335c255..517e1f4eb53f 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/SegmentConversionUtils.java +++ b/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 getSegmentNamesForTable(String tableNameWithType, URI public static void uploadSegment(Map configs, List
httpHeaders, List 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 configs, List
httpHeaders, + List 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 configs, List
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 configs, List
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; diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutorTest.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutorTest.java index 66a6b84cb10c..e583d1af4322 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutorTest.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutorTest.java @@ -20,13 +20,20 @@ import java.io.File; import java.net.URI; +import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; import org.apache.commons.io.FileUtils; +import org.apache.hc.core5.http.Header; +import org.apache.hc.core5.http.HttpHeaders; import org.apache.pinot.common.metadata.segment.SegmentZKMetadataCustomMapModifier; import org.apache.pinot.common.metrics.MinionMetrics; +import org.apache.pinot.common.utils.FileUploadDownloadClient; +import org.apache.pinot.common.utils.TarCompressionUtils; import org.apache.pinot.core.common.MinionConstants; import org.apache.pinot.core.minion.PinotTaskConfig; import org.apache.pinot.minion.MinionContext; @@ -34,7 +41,9 @@ import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; import org.apache.pinot.segment.local.utils.SegmentPushUtils; +import org.apache.pinot.segment.spi.V1Constants; import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; import org.apache.pinot.spi.config.instance.InstanceType; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; @@ -53,8 +62,8 @@ import org.testng.annotations.Test; -/// Tests the [BaseSingleSegmentConversionExecutor#executeTask] upload-failure handling: a segment-upload failure -/// must propagate so the task is marked failed (and retried) rather than being silently reported as successful. +/// Tests executeTask: upload failures propagate, the default METADATA push is unchanged, and the opt-in controller-copy +/// METADATA push is safe for a same-name segment refresh. public class BaseSingleSegmentConversionExecutorTest { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "BaseSingleSegmentConversionExecutorTest"); @@ -67,10 +76,13 @@ public class BaseSingleSegmentConversionExecutorTest { private static final String SEGMENT_NAME = "testSegment"; private static final String TASK_TYPE = "TestSingleSegmentConversionTask"; private static final String TASK_ID = "Task_" + TASK_TYPE + "_0"; - private static final long SEGMENT_CRC = 100L; + private static final String DOWNLOAD_URL = "http://unused/download"; + // A CRC that never matches the built segment, so the converted segment always counts as changed. + private static final long STALE_SEGMENT_CRC = 100L; private static final String D1 = "d1"; private File _segmentIndexDir; + private long _segmentCrc; @BeforeClass public void setUp() @@ -95,6 +107,7 @@ public void setUp() driver.init(config, new GenericRowRecordReader(rows)); driver.build(); _segmentIndexDir = new File(SEGMENT_DIR, SEGMENT_NAME); + _segmentCrc = Long.parseLong(new SegmentMetadataImpl(_segmentIndexDir).getCrc()); Assert.assertTrue(DATA_DIR.mkdirs()); MinionContext.getInstance().setDataDir(DATA_DIR); @@ -110,9 +123,9 @@ public void testExecuteTaskRethrowsWhenUploadFails() Mockito.anyString(), Mockito.anyString(), Mockito.anyString(), Mockito.any(File.class))) .thenThrow(new RuntimeException("simulated upload failure")); - TestSingleSegmentConversionExecutor executor = new TestSingleSegmentConversionExecutor(); + TestSingleSegmentConversionExecutor executor = new TestSingleSegmentConversionExecutor(STALE_SEGMENT_CRC); try { - executor.executeTask(createTaskConfig()); + executor.executeTask(createTaskConfig(STALE_SEGMENT_CRC)); Assert.fail("executeTask must rethrow when segment upload fails, not report success"); } catch (RuntimeException e) { Assert.assertEquals(e.getMessage(), "simulated upload failure"); @@ -125,8 +138,8 @@ public void testExecuteTaskSucceedsWhenUploadSucceeds() throws Exception { try (MockedStatic mocked = Mockito.mockStatic(SegmentConversionUtils.class)) { // uploadSegment is a no-op by default for the mocked static, simulating a successful upload. - TestSingleSegmentConversionExecutor executor = new TestSingleSegmentConversionExecutor(); - SegmentConversionResult result = executor.executeTask(createTaskConfig()); + TestSingleSegmentConversionExecutor executor = new TestSingleSegmentConversionExecutor(STALE_SEGMENT_CRC); + SegmentConversionResult result = executor.executeTask(createTaskConfig(STALE_SEGMENT_CRC)); Assert.assertEquals(result.getSegmentName(), SEGMENT_NAME); Assert.assertEquals(result.getTableNameWithType(), TABLE_NAME_WITH_TYPE); mocked.verify(() -> SegmentConversionUtils.uploadSegment(Mockito.any(), Mockito.any(), Mockito.any(), @@ -134,17 +147,6 @@ public void testExecuteTaskSucceedsWhenUploadSucceeds() } } - private PinotTaskConfig createTaskConfig() { - Map configs = new HashMap<>(); - configs.put(MinionConstants.TABLE_NAME_KEY, TABLE_NAME_WITH_TYPE); - configs.put(MinionConstants.SEGMENT_NAME_KEY, SEGMENT_NAME); - configs.put(MinionConstants.DOWNLOAD_URL_KEY, "http://unused/download"); - configs.put(MinionConstants.UPLOAD_URL_KEY, "http://unused/upload"); - configs.put(MinionConstants.ORIGINAL_SEGMENT_CRC_KEY, Long.toString(SEGMENT_CRC)); - configs.put("TASK_ID", TASK_ID); - return new PinotTaskConfig(TASK_TYPE, configs); - } - /// Verifies that when a METADATA-mode push fails after the converted tar was already staged to the output PinotFS, /// the staged tar is deleted before the exception propagates. Without this cleanup the rethrow would make the retry /// fail in moveSegmentToOutputPinotFS with "Output file already exists" (overwriteOutput defaults to false), so @@ -167,9 +169,9 @@ public void testExecuteTaskCleansUpStagedTarWhenMetadataPushFails() Mockito.any(), Mockito.anyList(), Mockito.anyList())) .thenThrow(new RuntimeException("simulated metadata push failure")); - TestSingleSegmentConversionExecutor executor = new TestSingleSegmentConversionExecutor(); + TestSingleSegmentConversionExecutor executor = new TestSingleSegmentConversionExecutor(STALE_SEGMENT_CRC); try { - executor.executeTask(createMetadataPushTaskConfig(outputDir)); + executor.executeTask(createMetadataPushTaskConfig(STALE_SEGMENT_CRC, outputDir)); Assert.fail("executeTask must rethrow when metadata push fails"); } catch (RuntimeException e) { Assert.assertEquals(e.getMessage(), "simulated metadata push failure"); @@ -179,19 +181,226 @@ public void testExecuteTaskCleansUpStagedTarWhenMetadataPushFails() } } - private PinotTaskConfig createMetadataPushTaskConfig(File outputDir) { + + /// The default METADATA push still registers the staged tar's URI and never takes the controller-copy path. + @Test + public void testDefaultMetadataPushRegistersStagedUri() + throws Exception { + File outputDir = new File(TEMP_DIR, "output-default"); + FileUtils.forceMkdir(outputDir); + try (MockedStatic segmentPushUtils = Mockito.mockStatic(SegmentPushUtils.class, + Mockito.CALLS_REAL_METHODS); + MockedStatic conversionUtils = Mockito.mockStatic(SegmentConversionUtils.class)) { + segmentPushUtils.when(() -> SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(), Mockito.any(), + Mockito.any(), Mockito.anyList(), Mockito.anyList())).thenAnswer(invocation -> null); + + new TestSingleSegmentConversionExecutor(STALE_SEGMENT_CRC).executeTask( + createMetadataPushTaskConfig(STALE_SEGMENT_CRC, outputDir)); + + String stagedTarName = SEGMENT_NAME + TarCompressionUtils.TAR_GZ_FILE_EXTENSION; + segmentPushUtils.verify(() -> SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(), Mockito.any(), + Mockito.argThat((Map uriToTar) -> uriToTar.size() == 1 + && uriToTar.values().iterator().next().endsWith(stagedTarName)), + Mockito.anyList(), Mockito.anyList())); + conversionUtils.verifyNoInteractions(); + } + // The staged tar is the segment's download URL in this mode, so it stays. + Assert.assertTrue(new File(outputDir, SEGMENT_NAME + TarCompressionUtils.TAR_GZ_FILE_EXTENSION).isFile()); + } + + /// Controller-copy push of a changed segment: task-unique staging name, TAR guards plus copy flag, a metadata-only + /// tar built locally, and the staged tar deleted once the push is done. + @Test + public void testMetadataPushStagesTarAndRegistersMetadata() + throws Exception { + File outputDir = new File(TEMP_DIR, "output-changed"); + FileUtils.forceMkdir(outputDir); + File stagedTar = new File(outputDir, SEGMENT_NAME + "." + TASK_ID + TarCompressionUtils.TAR_GZ_FILE_EXTENSION); + File capturedMetadataTar = new File(TEMP_DIR, "captured-metadata.tar.gz"); + List
capturedHeaders = new ArrayList<>(); + + try (MockedStatic mocked = Mockito.mockStatic(SegmentConversionUtils.class)) { + stubUploadSegmentMetadata(mocked, invocation -> { + // The controller copies from the staged tar while handling the request, so it must exist at this point. + Assert.assertTrue(stagedTar.isFile(), "staged tar must exist while the metadata push is in flight"); + assertTarHoldsSegment(stagedTar, new File(TEMP_DIR, "untar-staged")); + capturedHeaders.addAll(invocation.getArgument(1)); + FileUtils.copyFile(invocation.getArgument(6), capturedMetadataTar); + }); + + TestSingleSegmentConversionExecutor executor = new TestSingleSegmentConversionExecutor(STALE_SEGMENT_CRC, true); + SegmentConversionResult result = + executor.executeTask(createMetadataPushTaskConfig(STALE_SEGMENT_CRC, outputDir)); + Assert.assertEquals(result.getSegmentName(), SEGMENT_NAME); + mocked.verify(() -> SegmentConversionUtils.uploadSegment(Mockito.any(), Mockito.any(), Mockito.any(), + Mockito.anyString(), Mockito.anyString(), Mockito.anyString(), Mockito.any(File.class)), Mockito.never()); + } + + // Nothing references the staged tar after the controller copied it, so it must be gone. + Assert.assertEquals(outputDir.list().length, 0, "staged tar must be deleted after the push"); + + Assert.assertEquals(headerValue(capturedHeaders, HttpHeaders.IF_MATCH), String.valueOf(STALE_SEGMENT_CRC)); + Assert.assertEquals(headerValue(capturedHeaders, FileUploadDownloadClient.CustomHeaders.REFRESH_ONLY), "true"); + Assert.assertNotNull( + headerValue(capturedHeaders, FileUploadDownloadClient.CustomHeaders.SEGMENT_ZK_METADATA_CUSTOM_MAP_MODIFIER)); + Assert.assertEquals(headerValue(capturedHeaders, FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE), + FileUploadDownloadClient.FileUploadType.METADATA.toString()); + Assert.assertEquals( + headerValue(capturedHeaders, FileUploadDownloadClient.CustomHeaders.COPY_SEGMENT_TO_DEEP_STORE), "true"); + String downloadUri = headerValue(capturedHeaders, FileUploadDownloadClient.CustomHeaders.DOWNLOAD_URI); + Assert.assertEquals(new File(URI.create(downloadUri)), stagedTar); + + // The metadata tar carries only the two files the controller reads, and they describe the converted segment. + File untarredMetadataDir = TarCompressionUtils.untar(capturedMetadataTar, new File(TEMP_DIR, "untar-metadata")) + .get(0); + Set fileNames = + java.util.Arrays.stream(untarredMetadataDir.listFiles()).map(File::getName).collect(Collectors.toSet()); + Assert.assertEquals(fileNames, + Set.of(V1Constants.MetadataKeys.METADATA_FILE_NAME, V1Constants.SEGMENT_CREATION_META)); + SegmentMetadataImpl pushedMetadata = new SegmentMetadataImpl(untarredMetadataDir); + Assert.assertEquals(pushedMetadata.getName(), SEGMENT_NAME); + Assert.assertEquals(Long.parseLong(pushedMetadata.getCrc()), _segmentCrc); + } + + /// A plain-path output dir (local deep store) must reach the controller as a file URI. + @Test + public void testMetadataPushQualifiesSchemelessOutputDir() + throws Exception { + File outputDir = new File(TEMP_DIR, "output-schemeless"); + FileUtils.forceMkdir(outputDir); + File stagedTar = new File(outputDir, SEGMENT_NAME + "." + TASK_ID + TarCompressionUtils.TAR_GZ_FILE_EXTENSION); + List
capturedHeaders = new ArrayList<>(); + + try (MockedStatic mocked = Mockito.mockStatic(SegmentConversionUtils.class)) { + stubUploadSegmentMetadata(mocked, invocation -> { + Assert.assertTrue(stagedTar.isFile()); + capturedHeaders.addAll(invocation.getArgument(1)); + }); + PinotTaskConfig taskConfig = createTaskConfig(STALE_SEGMENT_CRC); + Map configs = taskConfig.getConfigs(); + configs.put(BatchConfigProperties.PUSH_MODE, BatchConfigProperties.SegmentPushType.METADATA.name()); + configs.put(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI, outputDir.getAbsolutePath()); + new TestSingleSegmentConversionExecutor(STALE_SEGMENT_CRC, true).executeTask(taskConfig); + } + + URI downloadUri = URI.create(headerValue(capturedHeaders, FileUploadDownloadClient.CustomHeaders.DOWNLOAD_URI)); + Assert.assertEquals(downloadUri.getScheme(), "file"); + Assert.assertEquals(new File(downloadUri), stagedTar); + Assert.assertEquals(outputDir.list().length, 0, "staged tar must be deleted after the push"); + } + + /// A failed push propagates and leaves no staged tar behind. + @Test + public void testMetadataPushDeletesStagedTarWhenPushFails() + throws Exception { + File outputDir = new File(TEMP_DIR, "output-failed"); + FileUtils.forceMkdir(outputDir); + + try (MockedStatic mocked = Mockito.mockStatic(SegmentConversionUtils.class)) { + mocked.when(() -> SegmentConversionUtils.uploadSegmentMetadata(Mockito.any(), Mockito.anyList(), + Mockito.anyList(), Mockito.anyString(), Mockito.anyString(), Mockito.anyString(), + Mockito.any(File.class))) + .thenThrow(new RuntimeException("simulated metadata push failure")); + + TestSingleSegmentConversionExecutor executor = new TestSingleSegmentConversionExecutor(STALE_SEGMENT_CRC, true); + try { + executor.executeTask(createMetadataPushTaskConfig(STALE_SEGMENT_CRC, outputDir)); + Assert.fail("executeTask must rethrow when metadata push fails"); + } catch (RuntimeException e) { + Assert.assertEquals(e.getMessage(), "simulated metadata push failure"); + } + } + Assert.assertEquals(outputDir.list().length, 0, "staged tar must be deleted after a failed push"); + } + + /// A retry reuses the staging name and overwrites an interrupted attempt's leftover instead of failing. + @Test + public void testMetadataPushOverwritesStagedTarLeftByPreviousAttempt() + throws Exception { + File outputDir = new File(TEMP_DIR, "output-leftover"); + FileUtils.forceMkdir(outputDir); + File stagedTar = new File(outputDir, SEGMENT_NAME + "." + TASK_ID + TarCompressionUtils.TAR_GZ_FILE_EXTENSION); + FileUtils.writeStringToFile(stagedTar, "leftover from an interrupted attempt", StandardCharsets.UTF_8); + + try (MockedStatic mocked = Mockito.mockStatic(SegmentConversionUtils.class)) { + stubUploadSegmentMetadata(mocked, + invocation -> assertTarHoldsSegment(stagedTar, new File(TEMP_DIR, "untar-leftover"))); + new TestSingleSegmentConversionExecutor(STALE_SEGMENT_CRC, true).executeTask( + createMetadataPushTaskConfig(STALE_SEGMENT_CRC, outputDir)); + } + Assert.assertEquals(outputDir.list().length, 0, "staged tar must be deleted after the push"); + } + + /// An unchanged segment (same CRC) is re-registered against its download URL without staging or copying. + @Test + public void testMetadataPushRegistersMetadataOnlyWhenSegmentUnchanged() + throws Exception { + File outputDir = new File(TEMP_DIR, "output-unchanged"); + FileUtils.forceMkdir(outputDir); + List
capturedHeaders = new ArrayList<>(); + + try (MockedStatic mocked = Mockito.mockStatic(SegmentConversionUtils.class)) { + stubUploadSegmentMetadata(mocked, invocation -> { + Assert.assertEquals(outputDir.list().length, 0, "an unchanged segment must not be staged"); + capturedHeaders.addAll(invocation.getArgument(1)); + }); + new TestSingleSegmentConversionExecutor(_segmentCrc, true).executeTask( + createMetadataPushTaskConfig(_segmentCrc, outputDir)); + } + + Assert.assertEquals(headerValue(capturedHeaders, HttpHeaders.IF_MATCH), String.valueOf(_segmentCrc)); + Assert.assertEquals(headerValue(capturedHeaders, FileUploadDownloadClient.CustomHeaders.REFRESH_ONLY), "true"); + Assert.assertEquals(headerValue(capturedHeaders, FileUploadDownloadClient.CustomHeaders.DOWNLOAD_URI), + DOWNLOAD_URL); + Assert.assertEquals( + headerValue(capturedHeaders, FileUploadDownloadClient.CustomHeaders.COPY_SEGMENT_TO_DEEP_STORE), "false"); + Assert.assertEquals(outputDir.list().length, 0); + } + + private interface MetadataPushCheck { + void check(org.mockito.invocation.InvocationOnMock invocation) + throws Exception; + } + + private static void stubUploadSegmentMetadata(MockedStatic mocked, MetadataPushCheck check) { + mocked.when(() -> SegmentConversionUtils.uploadSegmentMetadata(Mockito.any(), Mockito.anyList(), Mockito.anyList(), + Mockito.anyString(), Mockito.anyString(), Mockito.anyString(), Mockito.any(File.class))) + .thenAnswer(invocation -> { + check.check(invocation); + return null; + }); + } + + private void assertTarHoldsSegment(File tarFile, File untarDir) + throws Exception { + FileUtils.deleteDirectory(untarDir); + File untarredSegmentDir = TarCompressionUtils.untar(tarFile, untarDir).get(0); + Assert.assertEquals(new SegmentMetadataImpl(untarredSegmentDir).getName(), SEGMENT_NAME); + } + + private static String headerValue(List
headers, String name) { + return headers.stream().filter(header -> header.getName().equals(name)).map(Header::getValue).findFirst() + .orElse(null); + } + + private PinotTaskConfig createTaskConfig(long originalSegmentCrc) { Map configs = new HashMap<>(); configs.put(MinionConstants.TABLE_NAME_KEY, TABLE_NAME_WITH_TYPE); configs.put(MinionConstants.SEGMENT_NAME_KEY, SEGMENT_NAME); - configs.put(MinionConstants.DOWNLOAD_URL_KEY, "http://unused/download"); + configs.put(MinionConstants.DOWNLOAD_URL_KEY, DOWNLOAD_URL); configs.put(MinionConstants.UPLOAD_URL_KEY, "http://unused/upload"); - configs.put(MinionConstants.ORIGINAL_SEGMENT_CRC_KEY, Long.toString(SEGMENT_CRC)); + configs.put(MinionConstants.ORIGINAL_SEGMENT_CRC_KEY, Long.toString(originalSegmentCrc)); configs.put("TASK_ID", TASK_ID); - configs.put(BatchConfigProperties.PUSH_MODE, BatchConfigProperties.SegmentPushType.METADATA.name()); - configs.put(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI, outputDir.toURI().toString()); return new PinotTaskConfig(TASK_TYPE, configs); } + private PinotTaskConfig createMetadataPushTaskConfig(long originalSegmentCrc, File outputDir) { + PinotTaskConfig taskConfig = createTaskConfig(originalSegmentCrc); + taskConfig.getConfigs().put(BatchConfigProperties.PUSH_MODE, BatchConfigProperties.SegmentPushType.METADATA.name()); + taskConfig.getConfigs().put(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI, outputDir.toURI().toString()); + return taskConfig; + } + @AfterClass public void tearDown() throws Exception { @@ -201,9 +410,25 @@ public void tearDown() FileUtils.deleteDirectory(TEMP_DIR); } - /// Minimal concrete executor that stubs out the infrastructure-dependent hooks (download, CRC check, conversion, ZK - /// metadata modifier) so `executeTask` runs to the upload step without a server, controller, or deep store. + /// Stubs download, CRC check, conversion (a copy, so the CRC is unchanged) and the ZK modifier. private class TestSingleSegmentConversionExecutor extends BaseSingleSegmentConversionExecutor { + private final long _zkSegmentCrc; + private final boolean _copyToDeepStore; + + TestSingleSegmentConversionExecutor(long zkSegmentCrc) { + this(zkSegmentCrc, false); + } + + TestSingleSegmentConversionExecutor(long zkSegmentCrc, boolean copyToDeepStore) { + _zkSegmentCrc = zkSegmentCrc; + _copyToDeepStore = copyToDeepStore; + } + + @Override + protected boolean isCopyToDeepStoreForMetadataPush() { + return _copyToDeepStore; + } + @Override protected File downloadSegmentToLocalAndUntar(String tableNameWithType, String segmentName, String deepstoreURL, String taskType, File tempDataDir, String suffix) @@ -215,7 +440,7 @@ protected File downloadSegmentToLocalAndUntar(String tableNameWithType, String s @Override protected long getSegmentCrc(String tableNameWithType, String segmentName) { - return SEGMENT_CRC; + return _zkSegmentCrc; } @Override