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..dfdd8cfaab 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,8 @@ import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.tuple.ImmutablePair; import org.apache.commons.lang3.tuple.Pair; +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 +988,112 @@ 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 = 404, message = "Segment file not found at the download URL"), + @ApiResponse(code = 409, message = "Segment already exists"), + @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, + @Suspended final AsyncResponse asyncResponse) { + try { + asyncResponse.resume(registerSegmentInternal(segmentZKMetadataJson, tableName, tableType)); + } catch (Throwable t) { + asyncResponse.resume(t); + } + } + + private SuccessResponse registerSegmentInternal(String segmentZKMetadataJson, String tableName, String tableType) { + if (StringUtils.isEmpty(tableName)) { + throw new ControllerApplicationException(LOGGER, "tableName is required", Response.Status.BAD_REQUEST); + } + + SegmentZKMetadata segmentZKMetadata; + try { + ZNRecord znRecord = JsonUtils.stringToObject(segmentZKMetadataJson, ZNRecord.class); + segmentZKMetadata = new SegmentZKMetadata(znRecord); + } 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); + } + String downloadUrl = segmentZKMetadata.getDownloadUrl(); + if (StringUtils.isEmpty(downloadUrl)) { + 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); + + try { + LOGGER.info("Registering segment: {} for table: {}", segmentName, tableNameWithType); + + TableConfig tableConfig = _pinotHelixResourceManager.getTableConfig(tableNameWithType); + if (tableConfig == null) { + throw new ControllerApplicationException(LOGGER, "Failed to find table: " + tableNameWithType, + Response.Status.BAD_REQUEST); + } + + // Validate that the segment file exists in the deep store. + // The download URL may be a VIP URL (https://...) while the actual file resides on NFS (file://...). + // Resolve the actual deep store path using the controller's data directory. + URI dataDirURI = ControllerFilePathProvider.getInstance().getDataDirURI(); + String encodedSegmentName = URIUtils.encode(segmentName); + URI segmentFileURI = URIUtils.getUri(dataDirURI.toString(), rawTableName, encodedSegmentName); + PinotFS pinotFS = PinotFSFactory.create(dataDirURI.getScheme()); + if (!pinotFS.exists(segmentFileURI)) { + LOGGER.error("Segment file not found in deep store at: {} for segment: {} of table: {}", + segmentFileURI, segmentName, tableNameWithType); + throw new ControllerApplicationException(LOGGER, + String.format("Segment file not found in deep store at: %s", segmentFileURI), + Response.Status.NOT_FOUND); + } + + ZKOperator zkOperator = new ZKOperator(_pinotHelixResourceManager, _controllerConf, _controllerMetrics); + zkOperator.registerSegment(tableConfig, segmentZKMetadata); + + LOGGER.info("Successfully registered segment: {} for table: {}", segmentName, tableNameWithType); + return new SuccessResponse( + "Successfully registered segment: " + segmentName + " of table: " + tableNameWithType); + } catch (WebApplicationException e) { + throw e; + } catch (Exception e) { + _controllerMetrics.addMeteredGlobalValue(ControllerMeter.CONTROLLER_SEGMENT_UPLOAD_ERROR, 1L); + _controllerMetrics.addMeteredTableValue(tableName, ControllerMeter.CONTROLLER_TABLE_SEGMENT_UPLOAD_ERROR, 1L); + throw new ControllerApplicationException(LOGGER, + "Exception while registering segment: " + e.getMessage(), Response.Status.INTERNAL_SERVER_ERROR, e); + } + } + @POST @Path("segments/{tableName}/startReplaceSegments") @Authorize(targetType = TargetType.TABLE, paramName = "tableName", action = Actions.Table.REPLACE_SEGMENT) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/upload/ZKOperator.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/upload/ZKOperator.java index 94b3645cf3..de6cfcede9 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/upload/ZKOperator.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/upload/ZKOperator.java @@ -180,6 +180,64 @@ public void completeSegmentsOperations(TableConfig tableConfig, FileUploadType u processExistingSegments(tableConfig, uploadType, enableParallelPushProtection, headers, existingSegmentsList); } + /** + * Registers a segment directly from pre-built ZK metadata, without file upload or deep store copy. + * Only handles new segment creation. Throws CONFLICT if the segment already exists. + */ + public void registerSegment(TableConfig tableConfig, SegmentZKMetadata segmentZKMetadata) { + String tableNameWithType = tableConfig.getTableName(); + String segmentName = segmentZKMetadata.getSegmentName(); + + ZNRecord existingSegmentMetadataZNRecord = + _pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, segmentName); + if (existingSegmentMetadataZNRecord != null && shouldProcessAsNewSegment(tableNameWithType, segmentName, + existingSegmentMetadataZNRecord, false)) { + LOGGER.warn("Removing segment ZK metadata (recovering from previous registration failure) for table: {}, " + + "segment: {}", tableNameWithType, segmentName); + Preconditions.checkState(_pinotHelixResourceManager.removeSegmentZKMetadata(tableNameWithType, segmentName), + "Failed to remove segment ZK metadata for table: %s, segment: %s", tableNameWithType, segmentName); + existingSegmentMetadataZNRecord = null; + } + + if (existingSegmentMetadataZNRecord != null) { + LOGGER.error("Segment: {} already exists in table: {}, cannot register", segmentName, tableNameWithType); + throw new ControllerApplicationException(LOGGER, + String.format("Segment: %s already exists in table: %s.", segmentName, tableNameWithType), + Response.Status.CONFLICT); + } + + 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, true); + } catch (Exception e) { + LOGGER.error("Failed to assign segment: {} for table: {}, ZK metadata already cleaned up by assignSegment", + segmentName, tableNameWithType, e); + throw e; + } + + try { + // This call updates the sourceServer in the segment ZK metadata. + if (segmentZKMetadata.getSourceServer() != null) { + if (_pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata, 0)) { + LOGGER.info("Set sourceServer: {} for segment: {} of table: {}", + segmentZKMetadata.getSourceServer(), segmentName, tableNameWithType); + } else { + LOGGER.warn( + "Failed to persist sourceServer for segment: {} of table: {}, servers will fall back to deep store", + segmentName, tableNameWithType); + } + } + } catch (Exception e) { + LOGGER.error("Failed to update sourceServer. segment {}, table: {}", + segmentName, tableNameWithType, e); + } + } + public void completeReingestedSegmentOperations(String realtimeTableName, SegmentMetadata segmentMetadata, URI finalSegmentLocationURI, String sourceDownloadURIStr, String segmentDownloadURIStr, long segmentSizeInBytes) throws Exception { @@ -256,9 +314,10 @@ private void handleParallelPush(String tableNameWithType, String segmentName, lo } } - // TODO: peer download (sourceServer) is not supported for segment refresh because the refreshed segment - // is always fetched from deep store; enablePeerDownload only applies to new segment uploads via - // processNewSegment. Supporting peer download for refresh would require a separate design and implementation. + + // TODO: peer download (sourceServer) is not supported for segment refresh because the refreshed segment + // is always fetched from deep store; enablePeerDownload only applies to new segment uploads via + // processNewSegment. Supporting peer download for refresh would require a separate design and implementation. private void processExistingSegment(TableConfig tableConfig, SegmentMetadata segmentMetadata, FileUploadType uploadType, ZNRecord existingSegmentMetadataZNRecord, @Nullable URI finalSegmentLocationURI, File segmentFile, @Nullable String sourceDownloadURIStr, String segmentDownloadURIStr, @@ -273,6 +332,7 @@ private void processExistingSegment(TableConfig tableConfig, SegmentMetadata seg SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(existingSegmentMetadataZNRecord); long existingCrc = segmentZKMetadata.getCrc(); checkCRC(headers, tableNameWithType, segmentName, existingCrc); + segmentZKMetadata.setSourceServer(null); // Check segment upload start time when parallel push protection enabled if (enableParallelPushProtection) { @@ -370,6 +430,10 @@ private void processExistingSegment(TableConfig tableConfig, SegmentMetadata seg segmentDownloadURIStr, crypterName, segmentSizeInBytes); segmentZKMetadata.setCustomMap(customMapModifier.modifyMap(segmentZKMetadata.getCustomMap())); } + // Clear sourceServer to prevent peer download during refresh. Peer download from the source server + // during refresh causes a race condition: the source server may be replacing its segment directory + // (deleteDirectory + moveDirectory) while simultaneously serving peer download requests, producing + // corrupted tars with a mix of old and new segment files. if (!_pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata, expectedVersion)) { throw new RuntimeException( String.format("Failed to update ZK metadata for segment: %s, table: %s, expected version: %d", diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/api/upload/ZKOperatorTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/api/upload/ZKOperatorTest.java index e4e20a396d..487a00447c 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/api/upload/ZKOperatorTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/api/upload/ZKOperatorTest.java @@ -550,6 +550,45 @@ public void testExpectedVersionPreventsStaleWrite() 30_000L, "Failed to delete segment ZK metadata"); } + /** + * Tests that a CRC-mismatch refresh clears sourceServer from ZK metadata. During refresh, the source server + * replaces its segment directory (deleteDirectory + moveDirectory) while grizzly threads may be tarring + * the same directory for peer download, producing corrupted tars with a mix of old and new files. + */ + @Test + public void testCrcMismatchRefreshClearsSourceServer() + throws Exception { + String segmentName = "peerDownloadRefreshCrcMismatch"; + ZKOperator zkOperator = new ZKOperator(_resourceManager, mock(ControllerConf.class), mock(ControllerMetrics.class)); + + SegmentMetadata segmentMetadata = mock(SegmentMetadata.class); + when(segmentMetadata.getName()).thenReturn(segmentName); + when(segmentMetadata.getCrc()).thenReturn("11111"); + when(segmentMetadata.getIndexCreationTime()).thenReturn(1000L); + + // Initial upload with enablePeerDownload=true — sourceServer should be set + zkOperator.completeSegmentOperations(OFFLINE_TABLE_CONFIG, segmentMetadata, FileUploadType.SEGMENT, null, null, + "downloadUrl", "downloadUrl", null, 10, true, true, true, mock(HttpHeaders.class)); + assertNotNull(_resourceManager.getSegmentZKMetadata(OFFLINE_TABLE_NAME, segmentName).getSourceServer(), + "sourceServer should be set on initial upload"); + + // CRC-mismatch refresh — sourceServer must be cleared to prevent peer download race condition + when(segmentMetadata.getCrc()).thenReturn("22222"); + zkOperator.completeSegmentOperations(OFFLINE_TABLE_CONFIG, segmentMetadata, FileUploadType.SEGMENT, null, null, + "downloadUrl", "downloadUrl", null, 10, true, true, false, mock(HttpHeaders.class)); + + SegmentZKMetadata refreshedMetadata = _resourceManager.getSegmentZKMetadata(OFFLINE_TABLE_NAME, segmentName); + assertNotNull(refreshedMetadata); + assertNull(refreshedMetadata.getSourceServer(), + "sourceServer must be cleared on CRC-mismatch refresh to prevent peer download race condition"); + + // Cleanup + _resourceManager.deleteSegment(OFFLINE_TABLE_NAME, segmentName); + TestUtils.waitForCondition( + aVoid -> _resourceManager.getSegmentZKMetadata(OFFLINE_TABLE_NAME, segmentName) == null, + 30_000L, "Failed to delete segment ZK metadata"); + } + @AfterClass public void tearDown() { FileUtils.deleteQuietly(TEMP_DIR);