xiangfu0 commented on code in PR #19895:
URL: https://github.com/apache/pinot/pull/19895#discussion_r4289313801
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java:
##########
@@ -507,36 +541,225 @@ public static void
sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec, Pin
public static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec,
Map<String, File> segmentUriToMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters)
throws Exception {
+ sendSegmentsUriAndMetadata(spec, segmentUriToMetadataFileMap, headers,
parameters, 1);
+ }
+
+ @VisibleForTesting
+ static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec,
+ Map<String, File> segmentUriToMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters,
+ int parallelism)
+ throws Exception {
Map<String, File> segmentMetadataFileMap = new HashMap<>();
List<String> segmentURIs = new ArrayList<>();
for (Map.Entry<String, File> entry :
segmentUriToMetadataFileMap.entrySet()) {
String segmentName = getSegmentNameFromMetadataFile(entry.getValue());
- segmentMetadataFileMap.put(segmentName, entry.getValue());
+
Preconditions.checkArgument(segmentMetadataFileMap.putIfAbsent(segmentName,
entry.getValue()) == null,
+ "Multiple metadata files for segment: %s", segmentName);
segmentURIs.add(segmentName);
segmentURIs.add(entry.getKey());
}
- pushSegmentsMetadata(spec, segmentURIs, segmentMetadataFileMap, headers,
parameters);
+ pushSegmentsMetadata(spec, segmentURIs, segmentMetadataFileMap, headers,
parameters, parallelism);
}
private static void pushSegmentsMetadata(SegmentGenerationJobSpec spec,
Collection<String> segmentURIs,
- Map<String, File> segmentMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters)
+ Map<String, File> segmentMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters, int parallelism)
+ throws Exception {
+ List<MetadataBatch> batches = partitionMetadataBatches(segmentURIs,
segmentMetadataFileMap, parallelism);
+ validateMetadataControllers(spec);
Review Comment:
Consolidated controller validation in both batch input paths and removed it
from the shared push helper. The filesystem path still validates before
metadata generation.
Follow-up:
[80bb65ef12](https://github.com/apache/pinot/commit/80bb65ef12d460beac27e3e9b359bf3666d2fa15).
Validation: 27 focused tests and affected-module checks passed.
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java:
##########
@@ -507,36 +541,225 @@ public static void
sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec, Pin
public static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec,
Map<String, File> segmentUriToMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters)
throws Exception {
+ sendSegmentsUriAndMetadata(spec, segmentUriToMetadataFileMap, headers,
parameters, 1);
+ }
+
+ @VisibleForTesting
+ static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec,
+ Map<String, File> segmentUriToMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters,
+ int parallelism)
+ throws Exception {
Map<String, File> segmentMetadataFileMap = new HashMap<>();
List<String> segmentURIs = new ArrayList<>();
for (Map.Entry<String, File> entry :
segmentUriToMetadataFileMap.entrySet()) {
String segmentName = getSegmentNameFromMetadataFile(entry.getValue());
- segmentMetadataFileMap.put(segmentName, entry.getValue());
+
Preconditions.checkArgument(segmentMetadataFileMap.putIfAbsent(segmentName,
entry.getValue()) == null,
+ "Multiple metadata files for segment: %s", segmentName);
segmentURIs.add(segmentName);
segmentURIs.add(entry.getKey());
}
- pushSegmentsMetadata(spec, segmentURIs, segmentMetadataFileMap, headers,
parameters);
+ pushSegmentsMetadata(spec, segmentURIs, segmentMetadataFileMap, headers,
parameters, parallelism);
}
private static void pushSegmentsMetadata(SegmentGenerationJobSpec spec,
Collection<String> segmentURIs,
- Map<String, File> segmentMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters)
+ Map<String, File> segmentMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters, int parallelism)
+ throws Exception {
+ List<MetadataBatch> batches = partitionMetadataBatches(segmentURIs,
segmentMetadataFileMap, parallelism);
+ validateMetadataControllers(spec);
+ if (batches.isEmpty()) {
+ return;
+ }
+ if (batches.size() == 1) {
+ FileUploadDownloadClient client =
getOrCreateFileUploadDownloadClient(spec);
+ try {
+ pushMetadataBatch(spec, batches.get(0), headers, parameters, client,
() -> false);
+ } finally {
+ closeFileUploadDownloadClient(spec, client);
+ }
+ return;
+ }
+ pushMetadataBatches(spec, batches, headers, parameters);
+ }
+
+ /// Each batch owns its immutable name/URI pairs and metadata-file mapping;
source files stay caller-owned.
+ private static final class MetadataBatch {
+ private final List<String> _segmentUris;
+ private final Map<String, File> _metadataFiles;
+
+ private MetadataBatch(List<String> segmentUris, Map<String, File>
metadataFiles) {
+ _segmentUris = List.copyOf(segmentUris);
+ _metadataFiles = Map.copyOf(metadataFiles);
+ }
+ }
+
+ private static List<MetadataBatch>
partitionMetadataBatches(Collection<String> segmentURIs,
+ Map<String, File> metadataFiles, int parallelism) {
+ List<String> pairs = new ArrayList<>(segmentURIs);
+ Preconditions.checkArgument(pairs.size() == 2 * metadataFiles.size(),
"Incomplete metadata name/URI pairs");
+ Map<String, String> urisByName = new LinkedHashMap<>();
+ for (int i = 0; i < pairs.size(); i += 2) {
+ String name = pairs.get(i);
+ SegmentNameUtils.validatePartialOrFullSegmentName(name);
+ String uri = pairs.get(i + 1);
+ validateMetadataUris(List.of(uri));
+ Preconditions.checkArgument(urisByName.putIfAbsent(name, uri) == null,
+ "Multiple metadata files for segment: %s", name);
+ File metadata = metadataFiles.get(name);
+ Preconditions.checkArgument(metadata != null && metadata.isFile(),
"Missing metadata file for segment: %s", name);
+ }
+ int count = Math.min(Math.max(parallelism, 1), urisByName.size());
+ if (count == 0) {
+ return List.of();
+ }
+ List<MetadataBatch> batches = new ArrayList<>(count);
+ var entries = urisByName.entrySet().iterator();
+ int size = urisByName.size() / count;
+ int remainder = urisByName.size() % count;
+ for (int i = 0; i < count; i++) {
+ List<String> batchUris = new ArrayList<>();
+ Map<String, File> batchFiles = new LinkedHashMap<>();
+ for (int j = 0; j < size + (i < remainder ? 1 : 0); j++) {
+ Map.Entry<String, String> entry = entries.next();
+ batchUris.add(entry.getKey());
+ batchUris.add(entry.getValue());
+ batchFiles.put(entry.getKey(), metadataFiles.get(entry.getKey()));
+ }
+ batches.add(new MetadataBatch(batchUris, batchFiles));
+ }
+ return batches;
+ }
+
+ private static void validateMetadataUris(Collection<String> uris) {
+ for (String uri : uris) {
+ Preconditions.checkArgument(StringUtils.isNotBlank(uri) &&
!uri.contains("\n") && !uri.contains("\r"),
+ "Invalid metadata download URI: %s", uri);
+ URI.create(uri);
+ }
+ }
+
+ private static void validateMetadataControllers(SegmentGenerationJobSpec
spec)
+ throws URISyntaxException {
+ for (PinotClusterSpec cluster : spec.getPinotClusterSpecs()) {
+ URI controller = new URI(cluster.getControllerURI());
+ Preconditions.checkArgument(controller.getRawAuthority() != null && (
Review Comment:
This fail-fast check is intentional for batch metadata push: controllers
must be HTTP(S) URIs with an authority. It rejects invalid destinations before
metadata generation or upload; the non-batch path is unchanged.
Follow-up:
[80bb65ef12](https://github.com/apache/pinot/commit/80bb65ef12d460beac27e3e9b359bf3666d2fa15).
Validation: 27 focused tests and affected-module checks passed.
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java:
##########
@@ -507,36 +541,225 @@ public static void
sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec, Pin
public static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec,
Map<String, File> segmentUriToMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters)
throws Exception {
+ sendSegmentsUriAndMetadata(spec, segmentUriToMetadataFileMap, headers,
parameters, 1);
+ }
+
+ @VisibleForTesting
+ static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec,
+ Map<String, File> segmentUriToMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters,
+ int parallelism)
+ throws Exception {
Map<String, File> segmentMetadataFileMap = new HashMap<>();
List<String> segmentURIs = new ArrayList<>();
for (Map.Entry<String, File> entry :
segmentUriToMetadataFileMap.entrySet()) {
String segmentName = getSegmentNameFromMetadataFile(entry.getValue());
- segmentMetadataFileMap.put(segmentName, entry.getValue());
+
Preconditions.checkArgument(segmentMetadataFileMap.putIfAbsent(segmentName,
entry.getValue()) == null,
+ "Multiple metadata files for segment: %s", segmentName);
segmentURIs.add(segmentName);
segmentURIs.add(entry.getKey());
}
- pushSegmentsMetadata(spec, segmentURIs, segmentMetadataFileMap, headers,
parameters);
+ pushSegmentsMetadata(spec, segmentURIs, segmentMetadataFileMap, headers,
parameters, parallelism);
}
private static void pushSegmentsMetadata(SegmentGenerationJobSpec spec,
Collection<String> segmentURIs,
- Map<String, File> segmentMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters)
+ Map<String, File> segmentMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters, int parallelism)
+ throws Exception {
+ List<MetadataBatch> batches = partitionMetadataBatches(segmentURIs,
segmentMetadataFileMap, parallelism);
+ validateMetadataControllers(spec);
+ if (batches.isEmpty()) {
+ return;
+ }
+ if (batches.size() == 1) {
+ FileUploadDownloadClient client =
getOrCreateFileUploadDownloadClient(spec);
+ try {
+ pushMetadataBatch(spec, batches.get(0), headers, parameters, client,
() -> false);
+ } finally {
+ closeFileUploadDownloadClient(spec, client);
+ }
+ return;
+ }
+ pushMetadataBatches(spec, batches, headers, parameters);
+ }
+
+ /// Each batch owns its immutable name/URI pairs and metadata-file mapping;
source files stay caller-owned.
+ private static final class MetadataBatch {
+ private final List<String> _segmentUris;
+ private final Map<String, File> _metadataFiles;
+
+ private MetadataBatch(List<String> segmentUris, Map<String, File>
metadataFiles) {
+ _segmentUris = List.copyOf(segmentUris);
+ _metadataFiles = Map.copyOf(metadataFiles);
+ }
+ }
+
+ private static List<MetadataBatch>
partitionMetadataBatches(Collection<String> segmentURIs,
+ Map<String, File> metadataFiles, int parallelism) {
+ List<String> pairs = new ArrayList<>(segmentURIs);
+ Preconditions.checkArgument(pairs.size() == 2 * metadataFiles.size(),
"Incomplete metadata name/URI pairs");
+ Map<String, String> urisByName = new LinkedHashMap<>();
+ for (int i = 0; i < pairs.size(); i += 2) {
+ String name = pairs.get(i);
+ SegmentNameUtils.validatePartialOrFullSegmentName(name);
+ String uri = pairs.get(i + 1);
+ validateMetadataUris(List.of(uri));
Review Comment:
Download URIs are now validated once in each batch input path, with the
redundant partition-loop check removed. Keeping the filesystem check before
metadata generation preserves early failure; the existing invalid-input test
now covers both paths and verifies no filesystem or HTTP work begins.
Follow-up:
[80bb65ef12](https://github.com/apache/pinot/commit/80bb65ef12d460beac27e3e9b359bf3666d2fa15).
Validation: 27 focused tests and affected-module checks passed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
