xiangfu0 commented on code in PR #19895:
URL: https://github.com/apache/pinot/pull/19895#discussion_r4288750695
##########
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:
Minor: `validateMetadataControllers(spec)` runs here, but the
`Map<String,String>` entry path already calls it (and `validateMetadataUris`)
before `generateSegmentMetadataFiles`. The `Map<String,File>` entry path does
not, so the call is needed for that path but redundant for the other. Not a
correctness problem, just a little duplicated work per push — could be
consolidated into the two public entry points.
##########
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 now rejects any controller URI that is not http/https (or has no
authority) before any HTTP is attempted. That is reasonable for a batch
metadata push, but it is slightly stricter than the old path, which parsed the
URI and only failed later at the actual request. If someone had a cluster spec
that previously limped through for non-batch reasons, this will fail earlier.
Low risk and arguably an improvement — just worth being aware it is a new
fail-fast.
--
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]
