From 352aa30ac6e7e509e096c3f9a031d4bd938d6242 Mon Sep 17 00:00:00 2001 From: dinoocch Date: Thu, 12 Mar 2026 11:03:56 -0700 Subject: [PATCH] Add a register endpoint to upload zk metadata directly --- ...tSegmentUploadDownloadRestletResource.java | 132 ++++++++++++++++++ 1 file changed, 132 insertions(+) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java index e612901079..8216b0bdc5 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java @@ -71,6 +71,9 @@ import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.tuple.ImmutablePair; import org.apache.commons.lang3.tuple.Pair; +import org.apache.helix.model.IdealState; +import org.apache.helix.zookeeper.datamodel.ZNRecord; +import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; import org.apache.pinot.common.metrics.ControllerGauge; import org.apache.pinot.common.metrics.ControllerMeter; import org.apache.pinot.common.metrics.ControllerMetrics; @@ -986,6 +989,135 @@ public void uploadSegmentAsMultiPartV2(FormDataMultiPart multiPart, } } + @POST + @ManagedAsync + @Produces(MediaType.APPLICATION_JSON) + @Consumes(MediaType.APPLICATION_JSON) + @Path("/segments/register") + @Authorize(targetType = TargetType.TABLE, paramName = "tableName", action = Actions.Table.UPLOAD_SEGMENT) + @Authenticate(AccessType.CREATE) + @ApiOperation(value = "Register a segment from pre-built ZK metadata", + notes = "Registers a segment directly from its ZK metadata JSON without file upload, unpacking, or copying. " + + "The segment files must already exist at the downloadUrl in the ZK metadata. " + + "Request body is SegmentZKMetadata JSON: {\"segmentName\": \"...\", \"simpleFields\": {...}}. " + + "Required simpleFields: segment.download.url, segment.crc. " + + "Optional but recommended: segment.total.docs, segment.start.time, segment.end.time, " + + "segment.time.unit, segment.index.version, segment.creation.time.") + @ApiResponses(value = { + @ApiResponse(code = 200, message = "Successfully registered segment"), + @ApiResponse(code = 400, message = "Bad Request"), + @ApiResponse(code = 409, message = "Segment already exists and refresh not permitted"), + @ApiResponse(code = 500, message = "Internal error") + }) + @TrackInflightRequestMetrics + @TrackedByGauge(gauge = ControllerGauge.SEGMENT_UPLOADS_IN_PROGRESS) + public void registerSegment( + String segmentZKMetadataJson, + @ApiParam(value = "Name of the table", required = true) + @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_NAME) String tableName, + @ApiParam(value = "Type of the table") + @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_TYPE) @DefaultValue("OFFLINE") String tableType, + @ApiParam(value = "Whether to refresh if the segment already exists") + @QueryParam(FileUploadDownloadClient.QueryParameters.ALLOW_REFRESH) @DefaultValue("true") boolean allowRefresh, + @Suspended final AsyncResponse asyncResponse) { + try { + asyncResponse.resume(registerSegmentInternal(segmentZKMetadataJson, tableName, tableType, allowRefresh)); + } catch (Throwable t) { + asyncResponse.resume(t); + } + } + + private SuccessResponse registerSegmentInternal(String segmentZKMetadataJson, String tableName, String tableType, + boolean allowRefresh) { + if (StringUtils.isEmpty(tableName)) { + throw new ControllerApplicationException(LOGGER, "tableName is required", Response.Status.BAD_REQUEST); + } + + SegmentZKMetadata segmentZKMetadata; + try { + segmentZKMetadata = SegmentZKMetadata.fromJsonString(segmentZKMetadataJson); + } catch (Exception e) { + throw new ControllerApplicationException(LOGGER, + "Failed to parse segment ZK metadata: " + e.getMessage(), Response.Status.BAD_REQUEST); + } + + String segmentName = segmentZKMetadata.getSegmentName(); + if (StringUtils.isEmpty(segmentName)) { + throw new ControllerApplicationException(LOGGER, "segmentName is required in ZK metadata", + Response.Status.BAD_REQUEST); + } + if (StringUtils.isEmpty(segmentZKMetadata.getDownloadUrl())) { + throw new ControllerApplicationException(LOGGER, + "segment.download.url is required in simpleFields", Response.Status.BAD_REQUEST); + } + + String rawTableName = TableNameBuilder.extractRawTableName(tableName); + TableType type = TableType.valueOf(tableType.toUpperCase()); + String tableNameWithType = TableNameBuilder.forType(type).tableNameWithType(rawTableName); + + TableConfig tableConfig = _pinotHelixResourceManager.getTableConfig(tableNameWithType); + if (tableConfig == null) { + throw new ControllerApplicationException(LOGGER, "Failed to find table: " + tableNameWithType, + Response.Status.BAD_REQUEST); + } + + LOGGER.info("Registering segment: {} for table: {}", segmentName, tableNameWithType); + + ZNRecord existingZNRecord = _pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, segmentName); + + // If ZK metadata exists but the segment is absent from the ideal state, a previous registration failed + // partway through. Remove the stale ZK entry and treat this as a new segment. + if (existingZNRecord != null) { + IdealState idealState = _pinotHelixResourceManager.getTableIdealState(tableNameWithType); + Preconditions.checkState(idealState != null, "Failed to find ideal state for table: %s", tableNameWithType); + if (idealState.getInstanceStateMap(segmentName) == null) { + LOGGER.warn("Removing stale segment ZK metadata (recovering from previous registration failure) " + + "for table: {}, segment: {}", tableNameWithType, segmentName); + Preconditions.checkState(_pinotHelixResourceManager.removeSegmentZKMetadata(tableNameWithType, segmentName), + "Failed to remove stale segment ZK metadata for table: %s, segment: %s", tableNameWithType, segmentName); + existingZNRecord = null; + } + } + + if (existingZNRecord != null) { + if (!allowRefresh) { + throw new ControllerApplicationException(LOGGER, + String.format("Segment: %s already exists in table: %s. Refresh not permitted.", segmentName, + tableNameWithType), Response.Status.CONFLICT); + } + LOGGER.info("Segment: {} already exists in table: {}, refreshing", segmentName, tableNameWithType); + segmentZKMetadata.setRefreshTime(System.currentTimeMillis()); + if (!_pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata, + existingZNRecord.getVersion())) { + throw new RuntimeException( + String.format("Failed to update ZK metadata for segment: %s of table: %s", segmentName, + tableNameWithType)); + } + _pinotHelixResourceManager.sendSegmentRefreshMessage(tableNameWithType, segmentName, true, true); + } else { + if (_pinotHelixResourceManager.needTieredSegmentAssignment(tableConfig)) { + _pinotHelixResourceManager.updateSegmentTargetTier(tableNameWithType, segmentZKMetadata, + _pinotHelixResourceManager.getSortedTiers(tableConfig)); + } + segmentZKMetadata.setPushTime(System.currentTimeMillis()); + if (!_pinotHelixResourceManager.createSegmentZkMetadata(tableNameWithType, segmentZKMetadata)) { + throw new RuntimeException( + String.format("Failed to create ZK metadata for segment: %s of table: %s", segmentName, + tableNameWithType)); + } + try { + _pinotHelixResourceManager.assignSegment(tableConfig, segmentZKMetadata); + } catch (Exception e) { + LOGGER.error("Failed to assign segment: {} for table: {}, ZK metadata already cleaned up by assignSegment", + segmentName, tableNameWithType, e); + throw e; + } + } + + LOGGER.info("Successfully registered segment: {} for table: {}", segmentName, tableNameWithType); + return new SuccessResponse("Successfully registered segment: " + segmentName + " of table: " + tableNameWithType); + } + @POST @Path("segments/{tableName}/startReplaceSegments") @Authorize(targetType = TargetType.TABLE, paramName = "tableName", action = Actions.Table.REPLACE_SEGMENT)