From c0a33a75158e40952ee9d37a67827c0520a3c543 Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Mon, 7 Sep 2026 14:24:38 +0530 Subject: [PATCH 1/2] 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 From 3f20c400649fa13c6386722b7b7593092f44a999 Mon Sep 17 00:00:00 2001 From: Kartik Khare Date: Mon, 7 Sep 2026 15:12:02 +0530 Subject: [PATCH 2/2] Build METADATA push metadata from the local segment instead of re-downloading the staged tar Every minion METADATA push uploaded the segment tar to the output filesystem and then downloaded it back into java.io.tmpdir to extract metadata.properties and creation.meta, doubling the minion's transfer per segment and leaving a full segment copy in the temp dir during the push. SegmentPushUtils gains generateSegmentMetadataFile(File, File, String), which builds the two-file metadata tar from a local segment directory, plus sendSegmentUriAndMetadata / sendSegmentsUriAndMetadata overloads that push caller-supplied metadata tars. The existing PinotFS overloads keep their behavior and share one private push loop with the new ones. BaseSingleSegmentConversionExecutor builds the metadata tar for every METADATA push and the default path uses the new overload. BaseMultipleSegmentsConversionExecutor builds one metadata tar per output segment and uses the new overloads in both the per-segment and the batch mode. Wire behavior is unchanged. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Y375AgHYsh1YqNSsfvF1a8 --- ...aseMultipleSegmentsConversionExecutor.java | 67 ++-- .../BaseSingleSegmentConversionExecutor.java | 20 +- .../plugin/minion/tasks/BaseTaskExecutor.java | 19 -- ...ultipleSegmentsConversionExecutorTest.java | 177 +++++++++- ...seSingleSegmentConversionExecutorTest.java | 42 ++- .../segment/local/utils/SegmentPushUtils.java | 307 +++++++++++------- .../local/utils/SegmentPushUtilsTest.java | 112 +++++++ 7 files changed, 542 insertions(+), 202 deletions(-) diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutor.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutor.java index 052dd2df721b..b5df8a4306e6 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutor.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutor.java @@ -36,6 +36,7 @@ import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.stream.Collectors; +import javax.annotation.Nullable; import org.apache.commons.collections4.CollectionUtils; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.StringUtils; @@ -57,7 +58,6 @@ import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.config.table.TableType; -import org.apache.pinot.spi.filesystem.PinotFS; import org.apache.pinot.spi.ingestion.batch.BatchConfigProperties; import org.apache.pinot.spi.ingestion.batch.spec.PushJobSpec; import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec; @@ -214,6 +214,10 @@ public List executeTask(PinotTaskConfig pinotTaskConfig int numOutputSegments = segmentConversionResults.size(); List tarredSegmentFiles = new ArrayList<>(numOutputSegments); + // METADATA pushes send only metadata.properties and creation.meta, built here from the local converted segment + // so the staged tar is never downloaded back. Null entries for TAR push. + List segmentMetadataTarFiles = new ArrayList<>(numOutputSegments); + BatchConfigProperties.SegmentPushType pushType = getSegmentPushType(taskConfigs); int count = 1; for (SegmentConversionResult segmentConversionResult : segmentConversionResults) { File convertedSegmentDir = segmentConversionResult.getFile(); @@ -226,6 +230,9 @@ public List executeTask(PinotTaskConfig pinotTaskConfig segmentConversionResult.getSegmentName() + TarCompressionUtils.TAR_GZ_FILE_EXTENSION); TarCompressionUtils.createCompressedTarFile(convertedSegmentDir, convertedSegmentTarFile); tarredSegmentFiles.add(convertedSegmentTarFile); + segmentMetadataTarFiles.add(pushType == BatchConfigProperties.SegmentPushType.TAR ? null + : SegmentPushUtils.generateSegmentMetadataFile(convertedSegmentDir, convertedTarredSegmentDir, + segmentConversionResult.getSegmentName())); if (!FileUtils.deleteQuietly(convertedSegmentDir)) { LOGGER.warn("Failed to delete converted segment: {}", convertedSegmentDir.getAbsolutePath()); } @@ -249,13 +256,14 @@ public List executeTask(PinotTaskConfig pinotTaskConfig SegmentUploadContext segmentUploadContext = new SegmentUploadContext(pinotTaskConfig, segmentConversionResults); preUploadSegments(segmentUploadContext); - Map segmentUriToTarPathMap = new HashMap<>(); + Map segmentUriToMetadataFileMap = new HashMap<>(); PushJobSpec pushJobSpec = getPushJobSpec(taskConfigs); boolean batchSegmentUpload = pushJobSpec.isBatchSegmentUpload(); // Upload the tarred segments for (int i = 0; i < numOutputSegments; i++) { File convertedTarredSegmentFile = tarredSegmentFiles.get(i); + File segmentMetadataTarFile = segmentMetadataTarFiles.get(i); SegmentConversionResult segmentConversionResult = segmentConversionResults.get(i); String resultSegmentName = segmentConversionResult.getSegmentName(); _eventObserver.notifyProgress(_pinotTaskConfig, "Uploading segment: " + resultSegmentName + " (" + (i + 1) @@ -294,20 +302,24 @@ public List executeTask(PinotTaskConfig pinotTaskConfig } if (batchSegmentUpload) { - updateSegmentUriToTarPathMap(taskConfigs, outputSegmentTarURI, segmentConversionResult, - segmentUriToTarPathMap, pushJobSpec); + for (String segmentUri : getSegmentUris(taskConfigs, outputSegmentTarURI, segmentConversionResult, + pushJobSpec)) { + segmentUriToMetadataFileMap.put(segmentUri, segmentMetadataTarFile); + } } else { String rawTableName = TableNameBuilder.extractRawTableName(tableNameWithType); - pushSegment(rawTableName, taskConfigs, outputSegmentTarURI, httpHeaders, parameters, segmentConversionResult); + pushSegment(rawTableName, taskConfigs, outputSegmentTarURI, httpHeaders, parameters, segmentConversionResult, + segmentMetadataTarFile); if (!FileUtils.deleteQuietly(convertedTarredSegmentFile)) { LOGGER.warn("Failed to delete tarred converted segment: {}", convertedTarredSegmentFile.getAbsolutePath()); } + FileUtils.deleteQuietly(segmentMetadataTarFile); } } if (batchSegmentUpload) { try { - pushSegments(tableNameWithType, taskConfigs, pinotTaskConfig, segmentUriToTarPathMap, pushJobSpec, + pushSegments(tableNameWithType, taskConfigs, pinotTaskConfig, segmentUriToMetadataFileMap, pushJobSpec, authProvider, segmentConversionResults); } finally { for (File convertedTarredSegmentFile : tarredSegmentFiles) { @@ -316,6 +328,7 @@ public List executeTask(PinotTaskConfig pinotTaskConfig convertedTarredSegmentFile.getAbsolutePath()); } } + segmentMetadataTarFiles.forEach(FileUtils::deleteQuietly); } } @@ -410,21 +423,17 @@ private void downloadAndUntarSegment(String tableNameWithType, String taskType, } } + /// Download URIs the controller should register for a staged segment tar (prefix/suffix rules from the push spec). @VisibleForTesting - void updateSegmentUriToTarPathMap(Map taskConfigs, URI outputSegmentTarURI, - SegmentConversionResult segmentConversionResult, Map segmentUriToTarPathMap, - PushJobSpec pushJobSpec) { + Set getSegmentUris(Map taskConfigs, URI outputSegmentTarURI, + SegmentConversionResult segmentConversionResult, PushJobSpec pushJobSpec) { String segmentName = segmentConversionResult.getSegmentName(); if (!taskConfigs.containsKey(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI)) { throw new RuntimeException("Output dir URI missing for metadata push while processing segment: " + segmentName); } URI outputSegmentDirURI = URI.create(taskConfigs.get(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI)); - Map localSegmentUriToTarPathMap = - SegmentPushUtils.getSegmentUriToTarPathMap(outputSegmentDirURI, pushJobSpec, - new String[]{outputSegmentTarURI.toString()}); - if (!localSegmentUriToTarPathMap.isEmpty()) { - segmentUriToTarPathMap.putAll(localSegmentUriToTarPathMap); - } + return SegmentPushUtils.getSegmentUriToTarPathMap(outputSegmentDirURI, pushJobSpec, + new String[]{outputSegmentTarURI.toString()}).keySet(); } @VisibleForTesting @@ -436,23 +445,20 @@ List
getSegmentPushCommonHeaders(PinotTaskConfig pinotTaskConfig, AuthPr } private void pushSegments(String tableNameWithType, Map taskConfigs, PinotTaskConfig pinotTaskConfig, - Map segmentUriToTarPathMap, PushJobSpec pushJobSpec, - AuthProvider authProvider, List segmentConversionResults) + Map segmentUriToMetadataFileMap, PushJobSpec pushJobSpec, AuthProvider authProvider, + List segmentConversionResults) throws Exception { String tableName = TableNameBuilder.extractRawTableName(tableNameWithType); SegmentGenerationJobSpec spec = generateSegmentGenerationJobSpec(tableName, taskConfigs, pushJobSpec); List
headers = getSegmentPushCommonHeaders(pinotTaskConfig, authProvider, segmentConversionResults); List parameters = getSegmentPushCommonParams(tableNameWithType); - - URI outputSegmentDirURI = URI.create(taskConfigs.get(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI)); - try (PinotFS outputFileFS = MinionTaskUtils.getOutputPinotFS(taskConfigs, outputSegmentDirURI)) { - SegmentPushUtils.sendSegmentsUriAndMetadata(spec, outputFileFS, segmentUriToTarPathMap, headers, parameters); - } + SegmentPushUtils.sendSegmentsUriAndMetadata(spec, segmentUriToMetadataFileMap, headers, parameters); } private void pushSegment(String tableName, Map taskConfigs, URI outputSegmentTarURI, - List
headers, List parameters, SegmentConversionResult segmentConversionResult) + List
headers, List parameters, SegmentConversionResult segmentConversionResult, + @Nullable File segmentMetadataTarFile) throws Exception { BatchConfigProperties.SegmentPushType pushType = getSegmentPushType(taskConfigs); LOGGER.info("Trying to push Pinot segment with push mode {} from {}", pushType, outputSegmentTarURI); @@ -470,17 +476,12 @@ private void pushSegment(String tableName, Map taskConfigs, URI uploadURL, tarFile); break; case METADATA: - if (taskConfigs.containsKey(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI)) { - URI outputSegmentDirURI = URI.create(taskConfigs.get(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI)); - try (PinotFS outputFileFS = MinionTaskUtils.getOutputPinotFS(taskConfigs, outputSegmentDirURI)) { - Map segmentUriToTarPathMap = - SegmentPushUtils.getSegmentUriToTarPathMap(outputSegmentDirURI, pushJobSpec, - new String[]{outputSegmentTarURI.toString()}); - SegmentPushUtils.sendSegmentUriAndMetadata(spec, outputFileFS, segmentUriToTarPathMap, headers, parameters); - } - } else { - throw new RuntimeException("Output dir URI missing for metadata push"); + Map segmentUriToMetadataFileMap = new HashMap<>(); + for (String segmentUri : getSegmentUris(taskConfigs, outputSegmentTarURI, segmentConversionResult, + pushJobSpec)) { + segmentUriToMetadataFileMap.put(segmentUri, segmentMetadataTarFile); } + SegmentPushUtils.sendSegmentUriAndMetadata(spec, segmentUriToMetadataFileMap, headers, parameters); break; default: throw new UnsupportedOperationException("Unrecognized push mode - " + pushType); 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 2da32aa7ff71..950c82b27002 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 @@ -23,6 +23,7 @@ import java.io.File; import java.net.URI; import java.util.ArrayList; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.UUID; @@ -147,12 +148,16 @@ public SegmentConversionResult executeTask(PinotTaskConfig pinotTaskConfig) 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. + // METADATA pushes send only metadata.properties and creation.meta, built here from the local converted segment + // so the staged tar is never downloaded back. A controller-copy push of an unchanged segment (same CRC) skips + // the tar and the staging and only re-registers its metadata. File segmentMetadataTarFile = null; boolean reuseExistingSegment = false; + if (pushType == BatchConfigProperties.SegmentPushType.METADATA) { + segmentMetadataTarFile = + SegmentPushUtils.generateSegmentMetadataFile(convertedSegmentDir, tempDataDir, segmentName); + } if (copyToDeepStore) { - segmentMetadataTarFile = createSegmentMetadataTarFile(convertedSegmentDir, tempDataDir, segmentName); long convertedSegmentCrc = Long.parseLong(new SegmentMetadataImpl(convertedSegmentDir).getCrc()); reuseExistingSegment = convertedSegmentCrc == Long.parseLong(originalSegmentCrc) && StringUtils.isNotEmpty(downloadURL); @@ -224,7 +229,7 @@ public SegmentConversionResult executeTask(PinotTaskConfig pinotTaskConfig) uploadURL, downloadURL, convertedTarredSegmentFile, segmentMetadataTarFile); } else { uploadSegmentWithMetadata(configs, pinotTaskConfig, segmentConversionResult, authProvider, parameters, - tableNameWithType, convertedTarredSegmentFile); + tableNameWithType, convertedTarredSegmentFile, segmentMetadataTarFile); } break; default: @@ -262,7 +267,7 @@ protected boolean isCopyToDeepStoreForMetadataPush() { /// [BatchConfigProperties#PUSH_CONTROLLER_URI] in configs. private void uploadSegmentWithMetadata(Map configs, PinotTaskConfig pinotTaskConfig, SegmentConversionResult segmentConversionResult, AuthProvider authProvider, List parameters, - String tableNameWithType, File convertedTarredSegmentFile) + String tableNameWithType, File convertedTarredSegmentFile, File segmentMetadataTarFile) throws Exception { if (!configs.containsKey(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI)) { throw new RuntimeException("Output dir URI missing for metadata push. Set " @@ -282,9 +287,10 @@ private void uploadSegmentWithMetadata(Map configs, PinotTaskCon try (PinotFS outputFileFS = MinionTaskUtils.getOutputPinotFS(configs, outputSegmentDirURI)) { Map segmentUriToTarPathMap = SegmentPushUtils.getSegmentUriToTarPathMap(outputSegmentDirURI, pushJobSpec, new String[]{outputSegmentTarURI.toString()}); + Map segmentUriToMetadataFileMap = new HashMap<>(); + segmentUriToTarPathMap.keySet().forEach(uri -> segmentUriToMetadataFileMap.put(uri, segmentMetadataTarFile)); try { - SegmentPushUtils.sendSegmentUriAndMetadata(spec, outputFileFS, segmentUriToTarPathMap, metadataHeaders, - parameters); + SegmentPushUtils.sendSegmentUriAndMetadata(spec, segmentUriToMetadataFileMap, metadataHeaders, parameters); } catch (Exception e) { // The tar was already staged to the output PinotFS before this failure. If the task is retried, the next // moveSegmentToOutputPinotFS() would fail with "Output file already exists" (overwriteOutput defaults to 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 61fe7b746015..dce1317d5f1e 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,7 +20,6 @@ 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; @@ -45,14 +44,12 @@ 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; @@ -249,22 +246,6 @@ protected void deleteFromOutputPinotFS(Map configs, URI fileURI) } } - /// 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/test/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutorTest.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutorTest.java index 6793e1ad175c..7f1a2c15529b 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutorTest.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutorTest.java @@ -22,9 +22,12 @@ import java.io.IOException; import java.net.URI; import java.util.ArrayList; +import java.util.Arrays; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.UUID; import org.apache.commons.io.FileUtils; import org.apache.hc.core5.http.Header; @@ -32,17 +35,39 @@ import org.apache.hc.core5.http.message.BasicNameValuePair; import org.apache.pinot.common.auth.NullAuthProvider; 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.MinionConf; import org.apache.pinot.minion.MinionContext; +import org.apache.pinot.minion.event.MinionEventObservers; +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.auth.AuthProvider; +import org.apache.pinot.spi.config.instance.InstanceType; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +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.PushJobSpec; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.mockito.MockedStatic; import org.mockito.Mockito; +import org.mockito.stubbing.Answer; import org.testng.Assert; import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; @@ -237,7 +262,7 @@ public void testGetSegmentPushCommonParams() { } @Test - public void testUpdateSegmentURIToTarPathMap() + public void testGetSegmentUris() throws IOException { // setup File segmentDir = new File(_tempDir, "segments"); @@ -254,14 +279,154 @@ public void testUpdateSegmentURIToTarPathMap() PushJobSpec pushJobSpec = new PushJobSpec(); SegmentConversionResult conversionResult = new SegmentConversionResult.Builder().setSegmentName("mySegment").build(); - Map segmentUriToTarPathMap = new HashMap<>(); // test - _executor.updateSegmentUriToTarPathMap(taskConfigs, outputSegmentTarURI, conversionResult, segmentUriToTarPathMap, - pushJobSpec); + Set segmentUris = _executor.getSegmentUris(taskConfigs, outputSegmentTarURI, conversionResult, pushJobSpec); // validate - Assert.assertEquals(segmentUriToTarPathMap.size(), 1); - Assert.assertTrue(segmentUriToTarPathMap.containsKey(_tempDir.toURI() + "segments/segment.tar.gz")); + Assert.assertEquals(segmentUris, Set.of(_tempDir.toURI() + "segments/segment.tar.gz")); + } + + @DataProvider(name = "batchSegmentUpload") + public Object[][] batchSegmentUpload() { + return new Object[][]{{false}, {true}}; + } + + /// On METADATA push the metadata tar of every output segment is built from the local converted segment and pushed + /// with the local-metadata overload, so the staged tar is never downloaded back. Covers per-segment and batch pushes. + @Test(dataProvider = "batchSegmentUpload") + public void testMetadataPushUsesLocalMetadataTars(boolean batchSegmentUpload) + throws Exception { + File outputDir = new File(_tempDir, "output"); + FileUtils.forceMkdir(outputDir); + File dataDir = new File(_tempDir, "minionData"); + FileUtils.forceMkdir(dataDir); + File capturedDir = new File(_tempDir, "captured"); + FileUtils.forceMkdir(capturedDir); + File inputSegmentDir = buildSegment("input_0"); + List outputSegmentNames = List.of("merged_0", "merged_1"); + Map outputSegmentDirs = new HashMap<>(); + for (String name : outputSegmentNames) { + outputSegmentDirs.put(name, buildSegment(name)); + } + String taskId = "Task_TestMultiSegmentTask_" + UUID.randomUUID(); + MinionContext.getInstance().setDataDir(dataDir); + MinionMetrics.register(Mockito.mock(MinionMetrics.class)); + MinionEventObservers.getInstance().addMinionEventObserver(taskId, MinionTaskTestUtils.getMinionProgressObserver()); + + Map configs = new HashMap<>(); + configs.put(MinionConstants.TABLE_NAME_KEY, "myTable_OFFLINE"); + configs.put(MinionConstants.SEGMENT_NAME_KEY, "input_0"); + configs.put(MinionConstants.DOWNLOAD_URL_KEY, "http://unused/download"); + configs.put(MinionConstants.UPLOAD_URL_KEY, "http://unused/upload"); + configs.put(BatchConfigProperties.PUSH_MODE, BatchConfigProperties.SegmentPushType.METADATA.name()); + configs.put(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI, outputDir.toURI().toString()); + configs.put(BatchConfigProperties.BATCH_SEGMENT_UPLOAD, String.valueOf(batchSegmentUpload)); + configs.put("TASK_ID", taskId); + + Map capturedUriToMetadataFile = new HashMap<>(); + try (MockedStatic mocked = Mockito.mockStatic(SegmentPushUtils.class, + Mockito.CALLS_REAL_METHODS)) { + Answer capture = invocation -> { + Map uriToMetadataFile = invocation.getArgument(1); + for (Map.Entry entry : uriToMetadataFile.entrySet()) { + File copy = new File(capturedDir, entry.getValue().getName()); + FileUtils.copyFile(entry.getValue(), copy); + capturedUriToMetadataFile.put(entry.getKey(), copy); + } + return null; + }; + mocked.when(() -> SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(), Mockito.anyMap(), Mockito.anyList(), + Mockito.anyList())).thenAnswer(capture); + mocked.when(() -> SegmentPushUtils.sendSegmentsUriAndMetadata(Mockito.any(), Mockito.anyMap(), Mockito.anyList(), + Mockito.anyList())).thenAnswer(capture); + + new TestMultipleSegmentsConversionExecutor(inputSegmentDir, outputSegmentDirs).executeTask( + new PinotTaskConfig("TestMultiSegmentTask", configs)); + + // Never the PinotFS variants, which download the staged tar back to extract the metadata + mocked.verify(() -> SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(), Mockito.any(PinotFS.class), + Mockito.anyMap(), Mockito.anyList(), Mockito.anyList()), Mockito.never()); + mocked.verify(() -> SegmentPushUtils.sendSegmentsUriAndMetadata(Mockito.any(), Mockito.any(PinotFS.class), + Mockito.anyMap(), Mockito.anyList(), Mockito.anyList()), Mockito.never()); + } finally { + MinionEventObservers.getInstance().removeMinionEventObserver(taskId); + MinionContext.getInstance().setDataDir(null); + } + + Assert.assertEquals(capturedUriToMetadataFile.size(), 2); + for (String name : outputSegmentNames) { + File stagedTar = new File(outputDir, name + TarCompressionUtils.TAR_GZ_FILE_EXTENSION); + Assert.assertTrue(stagedTar.isFile(), "staged tar is the download URL and must stay"); + File metadataTar = capturedUriToMetadataFile.get(stagedTar.toURI().toString()); + Assert.assertNotNull(metadataTar, "metadata pushed for " + stagedTar.toURI()); + Assert.assertEquals(metadataTar.getName(), name + Constants.METADATA_TAR_GZ_FILE_EXT); + File untarred = TarCompressionUtils.untar(metadataTar, new File(_tempDir, "untar-" + name)).get(0); + Assert.assertEquals(new HashSet<>(Arrays.asList(untarred.list())), + Set.of(V1Constants.MetadataKeys.METADATA_FILE_NAME, V1Constants.SEGMENT_CREATION_META)); + Assert.assertEquals(new SegmentMetadataImpl(untarred).getName(), name); + } + } + + private File buildSegment(String segmentName) + throws Exception { + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("myTable").build(); + Schema schema = new Schema.SchemaBuilder().addSingleValueDimension("d1", FieldSpec.DataType.INT).build(); + GenericRow row = new GenericRow(); + row.putValue("d1", 1); + SegmentGeneratorConfig config = new SegmentGeneratorConfig(tableConfig, schema); + config.setInstanceType(InstanceType.MINION); + config.setOutDir(new File(_tempDir, "built").getPath()); + config.setSegmentName(segmentName); + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(config, new GenericRowRecordReader(List.of(row))); + driver.build(); + return new File(new File(_tempDir, "built"), segmentName); + } + + /// Stubs download, the segment-existence pre-check and conversion so executeTask runs to the push step offline. + private static class TestMultipleSegmentsConversionExecutor extends BaseMultipleSegmentsConversionExecutor { + private final File _inputSegmentDir; + private final Map _outputSegmentDirs; + + TestMultipleSegmentsConversionExecutor(File inputSegmentDir, Map outputSegmentDirs) { + super(new MinionConf()); + _inputSegmentDir = inputSegmentDir; + _outputSegmentDirs = outputSegmentDirs; + } + + @Override + protected void preProcess(PinotTaskConfig pinotTaskConfig) { + } + + @Override + protected File downloadSegmentToLocalAndUntar(String tableNameWithType, String segmentName, String deepstoreURL, + String taskType, File tempDataDir, String suffix) + throws Exception { + File indexDir = new File(tempDataDir, "input" + suffix); + FileUtils.copyDirectory(_inputSegmentDir, indexDir); + return indexDir; + } + + @Override + protected List convert(PinotTaskConfig pinotTaskConfig, List segmentDirs, + File workingDir) + throws Exception { + List results = new ArrayList<>(); + for (Map.Entry entry : _outputSegmentDirs.entrySet()) { + File convertedDir = new File(workingDir, entry.getKey()); + FileUtils.copyDirectory(entry.getValue(), convertedDir); + results.add(new SegmentConversionResult.Builder().setFile(convertedDir) + .setTableNameWithType(pinotTaskConfig.getConfigs().get(MinionConstants.TABLE_NAME_KEY)) + .setSegmentName(entry.getKey()).build()); + } + return results; + } + + @Override + protected SegmentZKMetadataCustomMapModifier getSegmentZKMetadataCustomMapModifier(PinotTaskConfig pinotTaskConfig, + SegmentConversionResult segmentConversionResult) { + return new SegmentZKMetadataCustomMapModifier(SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE, Map.of()); + } } } 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 e583d1af4322..e90f7e3203c4 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 @@ -165,8 +165,8 @@ public void testExecuteTaskCleansUpStagedTarWhenMetadataPushFails() Mockito.mockStatic(SegmentPushUtils.class, Mockito.CALLS_REAL_METHODS)) { minionTaskUtils.when(() -> MinionTaskUtils.getOutputPinotFS(Mockito.any(), Mockito.any())) .thenReturn(mockOutputFS); - segmentPushUtils.when(() -> SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(), Mockito.any(), - Mockito.any(), Mockito.anyList(), Mockito.anyList())) + segmentPushUtils.when(() -> SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(), Mockito.anyMap(), + Mockito.anyList(), Mockito.anyList())) .thenThrow(new RuntimeException("simulated metadata push failure")); TestSingleSegmentConversionExecutor executor = new TestSingleSegmentConversionExecutor(STALE_SEGMENT_CRC); @@ -182,30 +182,40 @@ public void testExecuteTaskCleansUpStagedTarWhenMetadataPushFails() } - /// The default METADATA push still registers the staged tar's URI and never takes the controller-copy path. + /// The default METADATA push still registers the staged tar's URI, with a metadata tar built from the local + /// segment instead of the staged tar being downloaded back, and never takes the controller-copy path. @Test public void testDefaultMetadataPushRegistersStagedUri() throws Exception { File outputDir = new File(TEMP_DIR, "output-default"); FileUtils.forceMkdir(outputDir); + File capturedMetadataTar = new File(TEMP_DIR, "captured-default-metadata.tar.gz"); + List capturedUris = new ArrayList<>(); 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); + segmentPushUtils.when(() -> SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(), Mockito.anyMap(), + Mockito.anyList(), Mockito.anyList())).thenAnswer(invocation -> { + Map uriToMetadataFile = invocation.getArgument(1); + capturedUris.addAll(uriToMetadataFile.keySet()); + FileUtils.copyFile(uriToMetadataFile.values().iterator().next(), capturedMetadataTar); + return 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())); + // No PinotFS-based push, which is the variant that downloads the staged tar back. + segmentPushUtils.verify(() -> SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(), + Mockito.any(PinotFS.class), Mockito.anyMap(), Mockito.anyList(), Mockito.anyList()), Mockito.never()); conversionUtils.verifyNoInteractions(); } + File stagedTar = new File(outputDir, SEGMENT_NAME + TarCompressionUtils.TAR_GZ_FILE_EXTENSION); + Assert.assertEquals(capturedUris.size(), 1); + Assert.assertEquals(new File(URI.create(capturedUris.get(0))), stagedTar); + assertMetadataTarDescribesSegment(capturedMetadataTar, new File(TEMP_DIR, "untar-default-metadata")); // 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()); + Assert.assertTrue(stagedTar.isFile()); } /// Controller-copy push of a changed segment: task-unique staging name, TAR guards plus copy flag, a metadata-only @@ -250,9 +260,13 @@ public void testMetadataPushStagesTarAndRegistersMetadata() 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); + assertMetadataTarDescribesSegment(capturedMetadataTar, new File(TEMP_DIR, "untar-metadata")); + } + + /// The metadata tar carries only the two files the controller reads, and they describe the converted segment. + private void assertMetadataTarDescribesSegment(File metadataTar, File untarDir) + throws Exception { + File untarredMetadataDir = TarCompressionUtils.untar(metadataTar, untarDir).get(0); Set fileNames = java.util.Arrays.stream(untarredMetadataDir.listFiles()).map(File::getName).collect(Collectors.toSet()); Assert.assertEquals(fileNames, diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java index cd8bcc80a068..beae6bb63806 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java @@ -58,6 +58,7 @@ import org.apache.pinot.common.utils.tls.TlsUtils; import org.apache.pinot.segment.local.constants.SegmentUploadConstants; import org.apache.pinot.segment.spi.V1Constants; +import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths; import org.apache.pinot.segment.spi.creator.name.SegmentNameUtils; import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.config.table.TableType; @@ -342,87 +343,17 @@ public static void sendSegmentUris(SegmentGenerationJobSpec spec, List s public static void sendSegmentUriAndMetadata(SegmentGenerationJobSpec spec, PinotFS fileSystem, Map segmentUriToTarPathMap, List
headers, List parameters) throws Exception { - String tableName = spec.getTableSpec().getTableName(); LOGGER.info("Start pushing segment metadata: {} to locations: {} for table {}", segmentUriToTarPathMap, - Arrays.toString(spec.getPinotClusterSpecs()), tableName); + Arrays.toString(spec.getPinotClusterSpecs()), spec.getTableSpec().getTableName()); FileUploadDownloadClient fileUploadDownloadClient = getOrCreateFileUploadDownloadClient(spec); - int socketTimeoutMs = getSocketTimeoutMs(spec); try { - for (String segmentUriPath : segmentUriToTarPathMap.keySet()) { - String tarFilePath = segmentUriToTarPathMap.get(segmentUriPath); - String fileName = new File(tarFilePath).getName(); - // segments stored in Pinot deep store do not have .tar.gz extension - String segmentName = fileName.endsWith(Constants.TAR_GZ_FILE_EXT) - ? fileName.substring(0, fileName.length() - Constants.TAR_GZ_FILE_EXT.length()) : fileName; - SegmentNameUtils.validatePartialOrFullSegmentName(segmentName); - File segmentMetadataFile; - // Check if there is a segment metadata tar gz file named `segmentName.metadata.tar.gz`, already in the remote - // directory. This is to avoid generating a new segment metadata tar gz file every time we push a segment, - // which requires downloading the entire segment tar gz file. - - URI metadataTarGzFilePath = generateSegmentMetadataURI(tarFilePath, segmentName); - LOGGER.info("Checking if metadata tar gz file {} exists", metadataTarGzFilePath); - if (spec.getPushJobSpec().isPreferMetadataTarGz() && fileSystem.exists(metadataTarGzFilePath)) { - segmentMetadataFile = new File(FileUtils.getTempDirectory(), - "segmentMetadata-" + UUID.randomUUID() + TarCompressionUtils.TAR_GZ_FILE_EXTENSION); - if (segmentMetadataFile.exists()) { - FileUtils.forceDelete(segmentMetadataFile); - } - fileSystem.copyToLocalFile(metadataTarGzFilePath, segmentMetadataFile); - } else { - segmentMetadataFile = generateSegmentMetadataFile(fileSystem, URI.create(tarFilePath)); - } + for (Map.Entry entry : segmentUriToTarPathMap.entrySet()) { + String tarFilePath = entry.getValue(); + String segmentName = getSegmentName(tarFilePath); + File segmentMetadataFile = getSegmentMetadataFile(spec, fileSystem, tarFilePath, segmentName); try { - for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) { - URI controllerURI; - try { - controllerURI = new URI(pinotClusterSpec.getControllerURI()); - } catch (URISyntaxException e) { - throw new RuntimeException("Got invalid controller uri - '" + pinotClusterSpec.getControllerURI() + "'"); - } - LOGGER.info("Pushing segment: {} to location: {} for table {}", segmentName, controllerURI, tableName); - int attempts = 1; - if (spec.getPushJobSpec() != null && spec.getPushJobSpec().getPushAttempts() > 0) { - attempts = spec.getPushJobSpec().getPushAttempts(); - } - long retryWaitMs = 1000L; - if (spec.getPushJobSpec() != null && spec.getPushJobSpec().getPushRetryIntervalMillis() > 0) { - retryWaitMs = spec.getPushJobSpec().getPushRetryIntervalMillis(); - } - RetryPolicies.exponentialBackoffRetryPolicy(attempts, retryWaitMs, 5).attempt(() -> { - List
reqHttpHeaders = new ArrayList<>(headers); - try { - reqHttpHeaders.add( - new BasicHeader(FileUploadDownloadClient.CustomHeaders.DOWNLOAD_URI, segmentUriPath)); - reqHttpHeaders.add(new BasicHeader(FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE, - FileUploadDownloadClient.FileUploadType.METADATA.toString())); - if (spec.getPushJobSpec() != null) { - reqHttpHeaders.add(new BasicHeader(FileUploadDownloadClient.CustomHeaders.COPY_SEGMENT_TO_DEEP_STORE, - String.valueOf(spec.getPushJobSpec().getCopyToDeepStoreForMetadataPush()))); - } - - SimpleHttpResponse response = fileUploadDownloadClient.uploadSegmentMetadata( - FileUploadDownloadClient.getUploadSegmentURI(controllerURI), segmentName, - segmentMetadataFile, reqHttpHeaders, parameters, socketTimeoutMs); - LOGGER.info("Response for pushing table {} segment {} to location {} - {}: {}", tableName, segmentName, - controllerURI, response.getStatusCode(), response.getResponse()); - return true; - } catch (HttpErrorStatusException e) { - int statusCode = e.getStatusCode(); - if (statusCode >= 500) { - // Temporary exception - LOGGER.warn("Caught temporary exception while pushing table: {} segment: {} to {}, will retry", - tableName, segmentName, controllerURI, e); - return false; - } else { - // Permanent exception - LOGGER.error("Caught permanent exception while pushing table: {} segment: {} to {}, won't retry", - tableName, segmentName, controllerURI, e); - throw e; - } - } - }); - } + pushSegmentMetadata(spec, fileUploadDownloadClient, entry.getKey(), segmentName, segmentMetadataFile, + headers, parameters); } finally { FileUtils.deleteQuietly(segmentMetadataFile); } @@ -432,30 +363,171 @@ public static void sendSegmentUriAndMetadata(SegmentGenerationJobSpec spec, Pino } } - public static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec, PinotFS fileSystem, - Map segmentUriToTarPathMap, List
headers, List parameters) + /// Same as above for segments whose metadata tars the caller already holds locally (see + /// [#generateSegmentMetadataFile(File, File, String)]), so nothing is downloaded back from the output filesystem. + /// Each file must be named `.metadata.tar.gz`, and the caller owns the files. + public static void sendSegmentUriAndMetadata(SegmentGenerationJobSpec spec, + Map segmentUriToMetadataFileMap, List
headers, List parameters) throws Exception { - String tableName = spec.getTableSpec().getTableName(); - ConcurrentHashMap segmentMetadataFileMap = new ConcurrentHashMap<>(); - ConcurrentLinkedQueue segmentURIs = new ConcurrentLinkedQueue<>(); - Map allSegmentsMetadataMap = new HashMap<>(); - File allSegmentsMetadataTarFile = null; - int nThreads = spec.getPushJobSpec().getSegmentMetadataGenerationParallelism(); + LOGGER.info("Start pushing local segment metadata for: {} to locations: {} for table {}", + segmentUriToMetadataFileMap.keySet(), Arrays.toString(spec.getPinotClusterSpecs()), + spec.getTableSpec().getTableName()); FileUploadDownloadClient fileUploadDownloadClient = getOrCreateFileUploadDownloadClient(spec); + try { + for (Map.Entry entry : segmentUriToMetadataFileMap.entrySet()) { + File segmentMetadataFile = entry.getValue(); + pushSegmentMetadata(spec, fileUploadDownloadClient, entry.getKey(), + getSegmentNameFromMetadataFile(segmentMetadataFile), segmentMetadataFile, headers, parameters); + } + } finally { + closeFileUploadDownloadClient(spec, fileUploadDownloadClient); + } + } + + private static void pushSegmentMetadata(SegmentGenerationJobSpec spec, + FileUploadDownloadClient fileUploadDownloadClient, String segmentUriPath, String segmentName, + File segmentMetadataFile, List
headers, List parameters) + throws Exception { + String tableName = spec.getTableSpec().getTableName(); int socketTimeoutMs = getSocketTimeoutMs(spec); - ExecutorService executor = Executors.newFixedThreadPool(nThreads); + for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) { + URI controllerURI; + try { + controllerURI = new URI(pinotClusterSpec.getControllerURI()); + } catch (URISyntaxException e) { + throw new RuntimeException("Got invalid controller uri - '" + pinotClusterSpec.getControllerURI() + "'"); + } + LOGGER.info("Pushing segment: {} to location: {} for table {}", segmentName, controllerURI, tableName); + int attempts = 1; + if (spec.getPushJobSpec() != null && spec.getPushJobSpec().getPushAttempts() > 0) { + attempts = spec.getPushJobSpec().getPushAttempts(); + } + long retryWaitMs = 1000L; + if (spec.getPushJobSpec() != null && spec.getPushJobSpec().getPushRetryIntervalMillis() > 0) { + retryWaitMs = spec.getPushJobSpec().getPushRetryIntervalMillis(); + } + RetryPolicies.exponentialBackoffRetryPolicy(attempts, retryWaitMs, 5).attempt(() -> { + List
reqHttpHeaders = new ArrayList<>(headers); + try { + reqHttpHeaders.add( + new BasicHeader(FileUploadDownloadClient.CustomHeaders.DOWNLOAD_URI, segmentUriPath)); + reqHttpHeaders.add(new BasicHeader(FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE, + FileUploadDownloadClient.FileUploadType.METADATA.toString())); + if (spec.getPushJobSpec() != null) { + reqHttpHeaders.add(new BasicHeader(FileUploadDownloadClient.CustomHeaders.COPY_SEGMENT_TO_DEEP_STORE, + String.valueOf(spec.getPushJobSpec().getCopyToDeepStoreForMetadataPush()))); + } + + SimpleHttpResponse response = fileUploadDownloadClient.uploadSegmentMetadata( + FileUploadDownloadClient.getUploadSegmentURI(controllerURI), segmentName, + segmentMetadataFile, reqHttpHeaders, parameters, socketTimeoutMs); + LOGGER.info("Response for pushing table {} segment {} to location {} - {}: {}", tableName, segmentName, + controllerURI, response.getStatusCode(), response.getResponse()); + return true; + } catch (HttpErrorStatusException e) { + int statusCode = e.getStatusCode(); + if (statusCode >= 500) { + // Temporary exception + LOGGER.warn("Caught temporary exception while pushing table: {} segment: {} to {}, will retry", + tableName, segmentName, controllerURI, e); + return false; + } else { + // Permanent exception + LOGGER.error("Caught permanent exception while pushing table: {} segment: {} to {}, won't retry", + tableName, segmentName, controllerURI, e); + throw e; + } + } + }); + } + } + + /// Metadata tar for a staged segment tar: the `.metadata.tar.gz` sidecar when the spec prefers it and + /// it exists, otherwise extracted from the segment tar (downloaded first unless the filesystem is local). + private static File getSegmentMetadataFile(SegmentGenerationJobSpec spec, PinotFS fileSystem, String tarFilePath, + String segmentName) + throws Exception { + URI metadataTarGzFilePath = generateSegmentMetadataURI(tarFilePath, segmentName); + LOGGER.info("Checking if metadata tar gz file {} exists", metadataTarGzFilePath); + if (spec.getPushJobSpec().isPreferMetadataTarGz() && fileSystem.exists(metadataTarGzFilePath)) { + File segmentMetadataFile = new File(FileUtils.getTempDirectory(), + SegmentUploadConstants.SEGMENT_METADATA_DIR_PREFIX + UUID.randomUUID() + + TarCompressionUtils.TAR_GZ_FILE_EXTENSION); + if (segmentMetadataFile.exists()) { + FileUtils.forceDelete(segmentMetadataFile); + } + fileSystem.copyToLocalFile(metadataTarGzFilePath, segmentMetadataFile); + return segmentMetadataFile; + } + return generateSegmentMetadataFile(fileSystem, URI.create(tarFilePath)); + } + + /// Segments stored in the deep store do not have the .tar.gz extension. + private static String getSegmentName(String tarFilePath) { + String fileName = new File(tarFilePath).getName(); + String segmentName = fileName.endsWith(Constants.TAR_GZ_FILE_EXT) + ? fileName.substring(0, fileName.length() - Constants.TAR_GZ_FILE_EXT.length()) : fileName; + SegmentNameUtils.validatePartialOrFullSegmentName(segmentName); + return segmentName; + } + + private static String getSegmentNameFromMetadataFile(File segmentMetadataFile) { + String fileName = segmentMetadataFile.getName(); + Preconditions.checkArgument(fileName.endsWith(Constants.METADATA_TAR_GZ_FILE_EXT), + "Segment metadata file: %s must be named %s", fileName, Constants.METADATA_TAR_GZ_FILE_EXT); + String segmentName = fileName.substring(0, fileName.length() - Constants.METADATA_TAR_GZ_FILE_EXT.length()); + SegmentNameUtils.validatePartialOrFullSegmentName(segmentName); + return segmentName; + } + + public static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec, PinotFS fileSystem, + Map segmentUriToTarPathMap, List
headers, List parameters) + throws Exception { LOGGER.info("Start pushing segment metadata: {} to locations: {} for table: {} with parallelism: {}", - segmentUriToTarPathMap, Arrays.toString(spec.getPinotClusterSpecs()), tableName, + segmentUriToTarPathMap, Arrays.toString(spec.getPinotClusterSpecs()), spec.getTableSpec().getTableName(), spec.getPushJobSpec().getPushParallelism()); - + ConcurrentHashMap segmentMetadataFileMap = new ConcurrentHashMap<>(); + ConcurrentLinkedQueue segmentURIs = new ConcurrentLinkedQueue<>(); + ExecutorService executor = + Executors.newFixedThreadPool(spec.getPushJobSpec().getSegmentMetadataGenerationParallelism()); try { generateSegmentMetadataFiles(spec, fileSystem, segmentUriToTarPathMap, segmentMetadataFileMap, segmentURIs, executor); - allSegmentsMetadataTarFile = createSegmentsMetadataTarFile(segmentURIs, segmentMetadataFileMap); - // the key is unused in batch upload mode and hence 'noopKey' - allSegmentsMetadataMap.put("noopKey", allSegmentsMetadataTarFile); + pushSegmentsMetadata(spec, segmentURIs, segmentMetadataFileMap, headers, parameters); + } finally { + for (File segmentMetadataFile : segmentMetadataFileMap.values()) { + FileUtils.deleteQuietly(segmentMetadataFile); + } + executor.shutdown(); + } + } - // perform metadata push in batch mode for every cluster + /// Batch variant of [#sendSegmentUriAndMetadata(SegmentGenerationJobSpec, Map, List, List)] for locally held + /// metadata tars named `.metadata.tar.gz`. The caller owns the files. + public static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec, + Map segmentUriToMetadataFileMap, List
headers, List parameters) + throws Exception { + Map segmentMetadataFileMap = new HashMap<>(); + List segmentURIs = new ArrayList<>(); + for (Map.Entry entry : segmentUriToMetadataFileMap.entrySet()) { + String segmentName = getSegmentNameFromMetadataFile(entry.getValue()); + segmentMetadataFileMap.put(segmentName, entry.getValue()); + segmentURIs.add(segmentName); + segmentURIs.add(entry.getKey()); + } + pushSegmentsMetadata(spec, segmentURIs, segmentMetadataFileMap, headers, parameters); + } + + private static void pushSegmentsMetadata(SegmentGenerationJobSpec spec, Collection segmentURIs, + Map segmentMetadataFileMap, List
headers, List parameters) + throws Exception { + String tableName = spec.getTableSpec().getTableName(); + FileUploadDownloadClient fileUploadDownloadClient = getOrCreateFileUploadDownloadClient(spec); + int socketTimeoutMs = getSocketTimeoutMs(spec); + File allSegmentsMetadataTarFile = createSegmentsMetadataTarFile(segmentURIs, segmentMetadataFileMap); + // the key is unused in batch upload mode and hence 'noopKey' + Map allSegmentsMetadataMap = Map.of("noopKey", allSegmentsMetadataTarFile); + try { for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) { URI controllerURI; try { @@ -500,14 +572,8 @@ public static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec, Pin }); } } finally { - for (Map.Entry metadataFileEntry : segmentMetadataFileMap.entrySet()) { - FileUtils.deleteQuietly(metadataFileEntry.getValue()); - } - if (allSegmentsMetadataTarFile != null) { - FileUtils.deleteQuietly(allSegmentsMetadataTarFile); - } + FileUtils.deleteQuietly(allSegmentsMetadataTarFile); closeFileUploadDownloadClient(spec, fileUploadDownloadClient); - executor.shutdown(); } } @@ -522,30 +588,8 @@ static void generateSegmentMetadataFiles(SegmentGenerationJobSpec spec, PinotFS futures.add( executor.submit(() -> { String tarFilePath = segmentUriToTarPathMap.get(segmentUriPath); - String fileName = new File(tarFilePath).getName(); - // segments stored in Pinot deep store do not have .tar.gz extension - String segmentName = fileName.endsWith(Constants.TAR_GZ_FILE_EXT) - ? fileName.substring(0, fileName.length() - Constants.TAR_GZ_FILE_EXT.length()) : fileName; - SegmentNameUtils.validatePartialOrFullSegmentName(segmentName); - File segmentMetadataFile; - // Check if there is a segment metadata tar gz file named `segmentName.metadata.tar.gz`, already in the - // remote directory. This is to avoid generating a new segment metadata tar gz file every time we push a - // segment, which requires downloading the entire segment tar gz file. - - URI metadataTarGzFilePath = generateSegmentMetadataURI(tarFilePath, segmentName); - LOGGER.info("Checking if metadata tar gz file {} exists", metadataTarGzFilePath); - if (spec.getPushJobSpec().isPreferMetadataTarGz() && fileSystem.exists(metadataTarGzFilePath)) { - segmentMetadataFile = new File(FileUtils.getTempDirectory(), - SegmentUploadConstants.SEGMENT_METADATA_DIR_PREFIX + UUID.randomUUID() - + TarCompressionUtils.TAR_GZ_FILE_EXTENSION); - if (segmentMetadataFile.exists()) { - FileUtils.forceDelete(segmentMetadataFile); - } - fileSystem.copyToLocalFile(metadataTarGzFilePath, segmentMetadataFile); - } else { - segmentMetadataFile = generateSegmentMetadataFile(fileSystem, URI.create(tarFilePath)); - } - segmentMetadataFileMap.put(segmentName, segmentMetadataFile); + String segmentName = getSegmentName(tarFilePath); + segmentMetadataFileMap.put(segmentName, getSegmentMetadataFile(spec, fileSystem, tarFilePath, segmentName)); segmentURIs.add(segmentName); segmentURIs.add(segmentUriPath); return null; @@ -707,6 +751,23 @@ public static File generateSegmentMetadataFile(PinotFS fileSystem, URI tarFileUR } } + /// Builds the metadata-only tar a METADATA push sends (`metadata.properties` and `creation.meta`) from a local + /// segment directory, named `.metadata.tar.gz`, so the segment tar never has to be downloaded back. + public static File generateSegmentMetadataFile(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); + } + } + public static URI generateSegmentMetadataURI(String segmentTarPath, String segmentName) throws URISyntaxException { URI segmentTarURI = URI.create(segmentTarPath); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/SegmentPushUtilsTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/SegmentPushUtilsTest.java index 0bf6005b8c30..a1763c472528 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/SegmentPushUtilsTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/SegmentPushUtilsTest.java @@ -34,13 +34,17 @@ import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.security.SecureRandom; +import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -58,6 +62,7 @@ import org.apache.pinot.common.utils.http.HttpClientConfig; import org.apache.pinot.common.utils.tls.TlsUtils; import org.apache.pinot.spi.filesystem.PinotFS; +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; @@ -375,6 +380,113 @@ private static void writeResponse(HttpExchange httpExchange, String responseBody } } + @Test + public void testGenerateSegmentMetadataFileFromLocalSegment() + throws Exception { + // A v3 segment layout with a data file that must not be included + File segmentDir = new File(_tempDir, TEST_SEGMENT_NAME); + File v3Dir = new File(segmentDir, "v3"); + FileUtils.forceMkdir(v3Dir); + FileUtils.writeStringToFile(new File(v3Dir, "metadata.properties"), "segment.name = " + TEST_SEGMENT_NAME, + StandardCharsets.UTF_8); + FileUtils.writeStringToFile(new File(v3Dir, "creation.meta"), "crc", StandardCharsets.UTF_8); + FileUtils.writeStringToFile(new File(v3Dir, "columns.psf"), "data", StandardCharsets.UTF_8); + + File metadataTar = SegmentPushUtils.generateSegmentMetadataFile(segmentDir, _tempDir, TEST_SEGMENT_NAME); + + assertEquals(metadataTar.getName(), TEST_SEGMENT_NAME + Constants.METADATA_TAR_GZ_FILE_EXT); + File untarred = TarCompressionUtils.untar(metadataTar, new File(_tempDir, "untar")).get(0); + assertEquals(new HashSet<>(Arrays.asList(untarred.list())), Set.of("metadata.properties", "creation.meta")); + assertEquals(FileUtils.readFileToString(new File(untarred, "metadata.properties"), StandardCharsets.UTF_8), + "segment.name = " + TEST_SEGMENT_NAME); + } + + @Test + public void testSendSegmentUriAndMetadataWithLocalMetadataFiles() + throws Exception { + List requests = new CopyOnWriteArrayList<>(); + HttpsServer httpsServer = createHttpsServer("/v2/segments", new RecordingHandler(requests)); + try { + URI controllerUri = new URI("https://localhost:" + httpsServer.getAddress().getPort()); + SegmentGenerationJobSpec spec = createSegmentGenerationJobSpec(controllerUri, createTlsSpec()); + spec.getPushJobSpec().setCopyToDeepStoreForMetadataPush(true); + File metadataTar = createMetadataTar(TEST_SEGMENT_NAME); + + SegmentPushUtils.sendSegmentUriAndMetadata(spec, Map.of(TEST_SEGMENT_URI, metadataTar), new ArrayList<>(), + new ArrayList<>()); + + assertEquals(requests.size(), 1); + Headers headers = requests.get(0); + assertEquals(headers.getFirst(FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE), + FileUploadDownloadClient.FileUploadType.METADATA.toString()); + assertEquals(headers.getFirst(FileUploadDownloadClient.CustomHeaders.DOWNLOAD_URI), TEST_SEGMENT_URI); + assertEquals(headers.getFirst(FileUploadDownloadClient.CustomHeaders.COPY_SEGMENT_TO_DEEP_STORE), "true"); + // The caller owns the metadata file + assertTrue(metadataTar.isFile()); + } finally { + httpsServer.stop(0); + } + } + + @Test + public void testSendSegmentsUriAndMetadataWithLocalMetadataFiles() + throws Exception { + List requests = new CopyOnWriteArrayList<>(); + HttpsServer httpsServer = createHttpsServer("/segments/batchUpload", new RecordingHandler(requests)); + try { + URI controllerUri = new URI("https://localhost:" + httpsServer.getAddress().getPort()); + SegmentGenerationJobSpec spec = createSegmentGenerationJobSpec(controllerUri, createTlsSpec()); + Map segmentUriToMetadataFileMap = new HashMap<>(); + segmentUriToMetadataFileMap.put("file:///tmp/segment1.tar.gz", createMetadataTar("segment1")); + segmentUriToMetadataFileMap.put("file:///tmp/segment2.tar.gz", createMetadataTar("segment2")); + + SegmentPushUtils.sendSegmentsUriAndMetadata(spec, segmentUriToMetadataFileMap, new ArrayList<>(), + new ArrayList<>()); + + // One batch request for both segments, and the caller still owns the metadata files + assertEquals(requests.size(), 1); + assertEquals(requests.get(0).getFirst(FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE), + FileUploadDownloadClient.FileUploadType.METADATA.toString()); + segmentUriToMetadataFileMap.values().forEach(file -> assertTrue(file.isFile())); + } finally { + httpsServer.stop(0); + } + } + + @Test + public void testSendSegmentUriAndMetadataRejectsUnnamedMetadataFile() { + SegmentGenerationJobSpec spec = createSegmentGenerationJobSpec(URI.create("https://localhost:1"), null); + File notAMetadataTar = new File(_tempDir, "segment.tar.gz"); + assertThrows(IllegalArgumentException.class, () -> SegmentPushUtils.sendSegmentUriAndMetadata(spec, + Map.of(TEST_SEGMENT_URI, notAMetadataTar), new ArrayList<>(), new ArrayList<>())); + } + + private File createMetadataTar(String segmentName) + throws IOException { + File segmentDir = new File(_tempDir, segmentName); + FileUtils.forceMkdir(segmentDir); + FileUtils.touch(new File(segmentDir, "metadata.properties")); + FileUtils.touch(new File(segmentDir, "creation.meta")); + return SegmentPushUtils.generateSegmentMetadataFile(segmentDir, _tempDir, segmentName); + } + + /// Records request headers and answers OK, for both single and batch metadata pushes. + private static class RecordingHandler implements HttpHandler { + private final List _requests; + + private RecordingHandler(List requests) { + _requests = requests; + } + + @Override + public void handle(HttpExchange httpExchange) + throws IOException { + _requests.add(httpExchange.getRequestHeaders()); + httpExchange.getRequestBody().readAllBytes(); + writeResponse(httpExchange, OK_RESPONSE); + } + } + @Test public void testGetOrCreateFileUploadDownloadClientUsesSharedDefaultClient() throws Exception {