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 51ad6d2ef936..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,10 +23,13 @@ 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; +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 +145,31 @@ public SegmentConversionResult executeTask(PinotTaskConfig pinotTaskConfig) reportTaskProcessingMetrics(tableNameWithType, taskType, segmentMetadata.getTotalDocs()); } + BatchConfigProperties.SegmentPushType pushType = getSegmentPushType(configs); + boolean copyToDeepStore = + pushType == BatchConfigProperties.SegmentPushType.METADATA && isCopyToDeepStoreForMetadataPush(); + // 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) { + 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 +204,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 +214,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 +224,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, segmentMetadataTarFile); + } break; default: throw new UnsupportedOperationException("Unrecognized push mode: " + pushType); @@ -212,9 +241,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,12 +256,18 @@ 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. 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 " @@ -249,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 @@ -268,6 +307,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..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 @@ -210,18 +210,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 +237,15 @@ 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); + } + } + /// 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/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 66a6b84cb10c..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 @@ -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 @@ -163,13 +165,13 @@ 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(); + 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,240 @@ public void testExecuteTaskCleansUpStagedTarWhenMetadataPushFails() } } - private PinotTaskConfig createMetadataPushTaskConfig(File outputDir) { + + /// 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.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)); + + // 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(stagedTar.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); + + 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, + 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 +424,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 +454,7 @@ protected File downloadSegmentToLocalAndUntar(String tableNameWithType, String s @Override protected long getSegmentCrc(String tableNameWithType, String segmentName) { - return SEGMENT_CRC; + return _zkSegmentCrc; } @Override 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 {