From e52e96d8b3beff918f3895505393e462e48a700e Mon Sep 17 00:00:00 2001 From: shuaiwan_LinkedIn Date: Wed, 18 Mar 2026 16:23:08 -0700 Subject: [PATCH] Implement SDS two-phase upload and segment distribution flow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the core Segment Distribution Service (SDS) protocol that enables P2P segment distribution without DeepStore dependency. Phase 1 - uploadSegmentIntent (controller): - Validates table/segment, creates ZK metadata with INTENT_CREATED status - Computes segment assignment, picks sourceServer, returns upload URLs Phase 2 - sdsUpload (server) + replicateSegment (controller): - Server receives segment bytes, writes to staging, untars, moves to final dir - Server calls controller to update IdealState (triggers OFFLINE→ONLINE) - sourceServer loads from disk via tryLoadExistingSegment() - Other replicas peer-download from sourceServer - Server waits for minimum replicas ONLINE before ACKing client Also includes: - SdsIntentCleanupTask: periodic cleanup of orphaned/stuck upload intents - Staging directory cleanup on server restart - Idempotency guard on replicateSegment for retry safety - Unit tests for SDS upload endpoint Co-Authored-By: Claude Opus 4.6 --- .../metadata/segment/SegmentZKMetadata.java | 8 + .../controller/BaseControllerStarter.java | 4 + ...tSegmentUploadDownloadRestletResource.java | 265 ++++++++++++ .../UploadSegmentIntentResponse.java | 71 ++++ .../helix/core/PinotHelixResourceManager.java | 26 ++ .../periodictask/SdsIntentCleanupTask.java | 103 +++++ ...ollerPeriodicTaskStarterStatelessTest.java | 2 +- .../SdsTwoPhaseUploadIntegrationTest.java | 392 ++++++++++++++++++ .../server/api/resources/TablesResource.java | 212 +++++++++- .../helix/HelixInstanceDataManager.java | 29 ++ .../pinot/server/api/BaseResourceTest.java | 1 + .../pinot/server/api/TablesResourceTest.java | 93 +++++ .../pinot/spi/utils/CommonConstants.java | 9 + 13 files changed, 1213 insertions(+), 2 deletions(-) create mode 100644 pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/UploadSegmentIntentResponse.java create mode 100644 pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/periodictask/SdsIntentCleanupTask.java create mode 100644 pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SdsTwoPhaseUploadIntegrationTest.java diff --git a/pinot-common/src/main/java/org/apache/pinot/common/metadata/segment/SegmentZKMetadata.java b/pinot-common/src/main/java/org/apache/pinot/common/metadata/segment/SegmentZKMetadata.java index 515d0f7fb6..314def04f5 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/metadata/segment/SegmentZKMetadata.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/metadata/segment/SegmentZKMetadata.java @@ -254,6 +254,14 @@ public void setSourceServer(String sourceServer) { setValue(Segment.SOURCE_SERVER, sourceServer); } + public String getSdsUploadStatus() { + return _simpleFields.get(Segment.SDS_UPLOAD_STATUS); + } + + public void setSdsUploadStatus(String status) { + setValue(Segment.SDS_UPLOAD_STATUS, status); + } + public String getCrypterName() { return _simpleFields.get(Segment.CRYPTER_NAME); } diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/BaseControllerStarter.java b/pinot-controller/src/main/java/org/apache/pinot/controller/BaseControllerStarter.java index 79e273ab8f..e1960f2a86 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/BaseControllerStarter.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/BaseControllerStarter.java @@ -100,6 +100,7 @@ import org.apache.pinot.controller.helix.core.minion.PinotHelixTaskResourceManager; import org.apache.pinot.controller.helix.core.minion.PinotTaskManager; import org.apache.pinot.controller.helix.core.minion.TaskMetricsEmitter; +import org.apache.pinot.controller.helix.core.periodictask.SdsIntentCleanupTask; import org.apache.pinot.controller.helix.core.realtime.PinotLLCRealtimeSegmentManager; import org.apache.pinot.controller.helix.core.realtime.SegmentCompletionConfig; import org.apache.pinot.controller.helix.core.realtime.SegmentCompletionManager; @@ -932,6 +933,9 @@ protected List setupControllerPeriodicTasks() { PeriodicTask resourceUtilizationChecker = new ResourceUtilizationChecker(_config, _connectionManager, _controllerMetrics, _diskUtilizationChecker, _executorService, _helixResourceManager); periodicTasks.add(resourceUtilizationChecker); + // Cleans up orphaned SDS upload intents (stale INTENT_CREATED / stuck UPLOAD_PROCESSING in ZK) + periodicTasks.add( + new SdsIntentCleanupTask(_helixResourceManager, _leadControllerManager, _controllerMetrics)); return periodicTasks; } 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..895f5285b5 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 @@ -46,6 +46,7 @@ import java.util.List; import java.util.Map; import java.util.UUID; +import java.util.concurrent.ThreadLocalRandom; import javax.annotation.Nullable; import javax.inject.Inject; import javax.ws.rs.Consumes; @@ -71,6 +72,11 @@ import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.tuple.ImmutablePair; import org.apache.commons.lang3.tuple.Pair; +import org.apache.helix.HelixAdmin; +import org.apache.helix.model.IdealState; +import org.apache.helix.model.InstanceConfig; +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; @@ -1303,4 +1309,263 @@ public static boolean validateMultiPart(Map> map, } return isGood; } + + // ======================== SDS (Segment Distribution Service) Endpoints ======================== + + @POST + @ManagedAsync + @Produces(MediaType.APPLICATION_JSON) + @Consumes(MediaType.APPLICATION_JSON) + @Path("/segments/{tableName}/uploadIntent") + @Authorize(targetType = TargetType.TABLE, paramName = "tableName", action = Actions.Table.UPLOAD_SEGMENT) + @Authenticate(AccessType.CREATE) + @ApiOperation(value = "Create an upload intent for a segment", + notes = "Phase 1 of the SDS 2-phase upload protocol. Validates the segment, creates ZK metadata, " + + "computes segment assignment, and returns the server URLs for direct upload. " + + "Request body is SegmentZKMetadata JSON.") + @ApiResponses(value = { + @ApiResponse(code = 200, message = "Successfully created upload intent"), + @ApiResponse(code = 400, message = "Bad Request"), + @ApiResponse(code = 409, message = "Segment already exists"), + @ApiResponse(code = 500, message = "Internal error") + }) + @TrackInflightRequestMetrics + @TrackedByGauge(gauge = ControllerGauge.SEGMENT_UPLOADS_IN_PROGRESS) + public void uploadSegmentIntent( + String segmentZKMetadataJson, + @ApiParam(value = "Name of the table", required = true) + @PathParam("tableName") String tableName, + @ApiParam(value = "Type of the table") + @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_TYPE) @DefaultValue("OFFLINE") String tableType, + @Suspended final AsyncResponse asyncResponse) { + try { + asyncResponse.resume(uploadSegmentIntentInternal(segmentZKMetadataJson, tableName, tableType)); + } catch (Throwable t) { + asyncResponse.resume(t); + } + } + + private UploadSegmentIntentResponse uploadSegmentIntentInternal( + String segmentZKMetadataJson, String tableName, String tableType) { + if (StringUtils.isEmpty(tableName)) { + throw new ControllerApplicationException(LOGGER, "tableName is required", Response.Status.BAD_REQUEST); + } + + // Parse segment metadata from 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 rawTableName = TableNameBuilder.extractRawTableName(tableName); + TableType type = TableType.valueOf(tableType.toUpperCase()); + String tableNameWithType = TableNameBuilder.forType(type).tableNameWithType(rawTableName); + + try { + LOGGER.info("Creating upload intent for segment: {} of table: {}", segmentName, tableNameWithType); + + TableConfig tableConfig = _pinotHelixResourceManager.getTableConfig(tableNameWithType); + if (tableConfig == null) { + throw new ControllerApplicationException(LOGGER, "Failed to find table: " + tableNameWithType, + Response.Status.BAD_REQUEST); + } + + // Check if segment already exists and is fully registered (in IS) + ZNRecord existingMetadata = + _pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, segmentName); + if (existingMetadata != null) { + IdealState idealState = _pinotHelixResourceManager.getTableIdealState(tableNameWithType); + if (idealState != null && idealState.getInstanceStateMap(segmentName) != null) { + throw new ControllerApplicationException(LOGGER, + String.format("Segment: %s already exists in table: %s", segmentName, tableNameWithType), + Response.Status.CONFLICT); + } + // ZK metadata exists but not in IS — previous failed attempt, clean up + LOGGER.warn("Removing stale segment ZK metadata for segment: {} of table: {}", + segmentName, tableNameWithType); + _pinotHelixResourceManager.removeSegmentZKMetadata(tableNameWithType, segmentName); + } + + // Create segment ZK metadata with INTENT_CREATED status. + // Set downloadUrl to empty string so non-source servers use peer-to-peer download + // (no deep store in SDS flow). pushTime is used by SdsIntentCleanupTask for staleness. + segmentZKMetadata.setPushTime(System.currentTimeMillis()); + segmentZKMetadata.setSdsUploadStatus(CommonConstants.Segment.SdsUploadStatus.INTENT_CREATED); + segmentZKMetadata.setDownloadUrl(CommonConstants.Segment.METADATA_URI_FOR_PEER_DOWNLOAD); + + if (!_pinotHelixResourceManager.createSegmentZkMetadata(tableNameWithType, segmentZKMetadata)) { + throw new RuntimeException( + String.format("Failed to create ZK metadata for segment: %s of table: %s", + segmentName, tableNameWithType)); + } + + // Compute segment assignment without updating IdealState yet — IS is updated later + // in replicateSegment() after the client uploads the segment bytes to the sourceServer. + List assignedInstances = + _pinotHelixResourceManager.computeSegmentAssignment(tableConfig, segmentName); + if (assignedInstances.isEmpty()) { + // Clean up ZK metadata + _pinotHelixResourceManager.removeSegmentZKMetadata(tableNameWithType, segmentName); + throw new ControllerApplicationException(LOGGER, + "No instances available for segment assignment", Response.Status.INTERNAL_SERVER_ERROR); + } + + // Pick sourceServer randomly from assigned instances + String sourceServer = + assignedInstances.get(ThreadLocalRandom.current().nextInt(assignedInstances.size())); + + // Persist sourceServer in ZK metadata + segmentZKMetadata.setSourceServer(sourceServer); + _pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata, 0); + + String uploadUrl = buildSourceServerUploadUrl(tableNameWithType, segmentName, sourceServer); + + LOGGER.info("Created upload intent for segment: {} of table: {}, sourceServer: {}, assignedInstances: {}", + segmentName, tableNameWithType, sourceServer, assignedInstances); + + return new UploadSegmentIntentResponse(segmentName, tableNameWithType, assignedInstances, sourceServer, + uploadUrl); + } catch (WebApplicationException e) { + throw e; + } catch (Exception e) { + _controllerMetrics.addMeteredGlobalValue(ControllerMeter.CONTROLLER_SEGMENT_UPLOAD_ERROR, 1L); + throw new ControllerApplicationException(LOGGER, + "Exception while creating upload intent: " + e.getMessage(), Response.Status.INTERNAL_SERVER_ERROR, e); + } + } + + private String buildSourceServerUploadUrl(String tableNameWithType, String segmentName, + String sourceServerInstanceId) { + HelixAdmin helixAdmin = _pinotHelixResourceManager.getHelixAdmin(); + String clusterName = _pinotHelixResourceManager.getHelixClusterName(); + InstanceConfig instanceConfig = helixAdmin.getInstanceConfig(clusterName, sourceServerInstanceId); + + String hostName = instanceConfig.getHostName(); + int httpsPort = instanceConfig.getRecord().getIntField( + CommonConstants.Helix.Instance.ADMIN_HTTPS_PORT_KEY, -1); + boolean useHttps = httpsPort > 0; + String scheme = useHttps ? CommonConstants.HTTPS_PROTOCOL : CommonConstants.HTTP_PROTOCOL; + int port = useHttps ? httpsPort : instanceConfig.getRecord().getIntField( + CommonConstants.Helix.Instance.ADMIN_PORT_KEY, CommonConstants.Server.DEFAULT_ADMIN_API_PORT); + + return scheme + "://" + hostName + ":" + port + "/segments/" + tableNameWithType + "/" + + URIUtils.encode(segmentName) + "/sdsUpload"; + } + + @POST + @ManagedAsync + @Produces(MediaType.APPLICATION_JSON) + @Path("/segments/{tableName}/{segmentName}/replicate") + @Authorize(targetType = TargetType.TABLE, paramName = "tableName", action = Actions.Table.UPLOAD_SEGMENT) + @Authenticate(AccessType.CREATE) + @ApiOperation(value = "Replicate a segment to all assigned servers", + notes = "Phase 2 callback from SDS. Updates IdealState to trigger OFFLINE to ONLINE transitions " + + "on all assigned servers.") + @ApiResponses(value = { + @ApiResponse(code = 200, message = "Successfully initiated replication"), + @ApiResponse(code = 400, message = "Bad Request"), + @ApiResponse(code = 404, message = "Segment metadata not found"), + @ApiResponse(code = 500, message = "Internal error") + }) + @TrackInflightRequestMetrics + public void replicateSegment( + @ApiParam(value = "Name of the table", required = true) + @PathParam("tableName") String tableName, + @ApiParam(value = "Name of the segment", required = true) + @PathParam("segmentName") String segmentName, + @ApiParam(value = "Type of the table") + @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_TYPE) @DefaultValue("OFFLINE") String tableType, + @Suspended final AsyncResponse asyncResponse) { + try { + asyncResponse.resume(replicateSegmentInternal(tableName, segmentName, tableType)); + } catch (Throwable t) { + asyncResponse.resume(t); + } + } + + private SuccessResponse replicateSegmentInternal(String tableName, String segmentName, String tableType) { + if (StringUtils.isEmpty(tableName) || StringUtils.isEmpty(segmentName)) { + throw new ControllerApplicationException(LOGGER, "tableName and segmentName are required", + 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("Replicating 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); + } + + // Verify ZK metadata exists with INTENT_CREATED status + ZNRecord zkRecord = + _pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, segmentName); + if (zkRecord == null) { + throw new ControllerApplicationException(LOGGER, + String.format("Segment ZK metadata not found for segment: %s of table: %s", + segmentName, tableNameWithType), + Response.Status.NOT_FOUND); + } + + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(zkRecord); + String sdsStatus = segmentZKMetadata.getSdsUploadStatus(); + + // Idempotency guard: the server may retry this call if it timed out on the first attempt. + // If the segment is already in IdealState, the previous call succeeded — just ensure status is correct. + IdealState idealState = _pinotHelixResourceManager.getTableIdealState(tableNameWithType); + if (idealState != null && idealState.getInstanceStateMap(segmentName) != null) { + LOGGER.info("Segment: {} already in IdealState for table: {}, skipping assignment (idempotent retry)", + segmentName, tableNameWithType); + // Ensure status is UPLOAD_COMPLETE + if (!CommonConstants.Segment.SdsUploadStatus.UPLOAD_COMPLETE.equals(sdsStatus)) { + segmentZKMetadata.setSdsUploadStatus(CommonConstants.Segment.SdsUploadStatus.UPLOAD_COMPLETE); + _pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata); + } + } else { + if (!CommonConstants.Segment.SdsUploadStatus.INTENT_CREATED.equals(sdsStatus)) { + LOGGER.warn("Unexpected SDS status: {} for segment: {} of table: {}, proceeding anyway", + sdsStatus, segmentName, tableNameWithType); + } + + // Update status to UPLOAD_PROCESSING + segmentZKMetadata.setSdsUploadStatus(CommonConstants.Segment.SdsUploadStatus.UPLOAD_PROCESSING); + _pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata); + + // Assign segment in IdealState. This triggers OFFLINE→ONLINE on all assigned servers. + // The sourceServer finds the segment already on disk (placed by sdsUpload) via tryLoadExistingSegment(). + // Other servers peer-download from the sourceServer (already ONLINE) via downloadSegmentFromPeers(). + _pinotHelixResourceManager.assignSegment(tableConfig, segmentZKMetadata, true); + + // Update status to UPLOAD_COMPLETE + segmentZKMetadata.setSdsUploadStatus(CommonConstants.Segment.SdsUploadStatus.UPLOAD_COMPLETE); + _pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata); + } + + LOGGER.info("Successfully initiated replication for segment: {} of table: {}", + segmentName, tableNameWithType); + return new SuccessResponse( + "Successfully initiated replication for segment: " + segmentName + " of table: " + tableNameWithType); + } catch (WebApplicationException e) { + throw e; + } catch (Exception e) { + _controllerMetrics.addMeteredGlobalValue(ControllerMeter.CONTROLLER_SEGMENT_UPLOAD_ERROR, 1L); + throw new ControllerApplicationException(LOGGER, + "Exception while replicating segment: " + e.getMessage(), Response.Status.INTERNAL_SERVER_ERROR, e); + } + } } diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/UploadSegmentIntentResponse.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/UploadSegmentIntentResponse.java new file mode 100644 index 0000000000..af27ca702e --- /dev/null +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/UploadSegmentIntentResponse.java @@ -0,0 +1,71 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.controller.api.resources; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import java.util.List; + + +public class UploadSegmentIntentResponse { + private final String _segmentName; + private final String _tableNameWithType; + private final List _assignedInstances; + private final String _sourceServer; + private final String _sourceServerUploadUrl; + + @JsonCreator + public UploadSegmentIntentResponse( + @JsonProperty("segmentName") String segmentName, + @JsonProperty("tableNameWithType") String tableNameWithType, + @JsonProperty("assignedInstances") List assignedInstances, + @JsonProperty("sourceServer") String sourceServer, + @JsonProperty("sourceServerUploadUrl") String sourceServerUploadUrl) { + _segmentName = segmentName; + _tableNameWithType = tableNameWithType; + _assignedInstances = assignedInstances; + _sourceServer = sourceServer; + _sourceServerUploadUrl = sourceServerUploadUrl; + } + + @JsonProperty("segmentName") + public String getSegmentName() { + return _segmentName; + } + + @JsonProperty("tableNameWithType") + public String getTableNameWithType() { + return _tableNameWithType; + } + + @JsonProperty("assignedInstances") + public List getAssignedInstances() { + return _assignedInstances; + } + + @JsonProperty("sourceServer") + public String getSourceServer() { + return _sourceServer; + } + + @JsonProperty("sourceServerUploadUrl") + public String getSourceServerUploadUrl() { + return _sourceServerUploadUrl; + } +} diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java index 67a8dc17f4..e58f655fd1 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java @@ -2597,6 +2597,32 @@ public void assignSegment(TableConfig tableConfig, SegmentZKMetadata segmentZKMe } } + /** + * Computes the segment assignment (target instances) without updating the IdealState. + * Used by uploadSegmentIntent() to determine which servers should host the segment. + * + * @return list of instance names that should host the segment + */ + public List computeSegmentAssignment(TableConfig tableConfig, String segmentName) { + String tableNameWithType = tableConfig.getTableName(); + + Map instancePartitionsMap; + if (TableNameBuilder.isOfflineTableResource(tableNameWithType)) { + instancePartitionsMap = getInstacePartitionsMap(tableConfig, null); + } else { + instancePartitionsMap = fetchOrComputeInstancePartitions(tableNameWithType, tableConfig); + } + + SegmentAssignment segmentAssignment = + SegmentAssignmentFactory.getSegmentAssignment(_helixZkManager, tableConfig, _controllerMetrics); + + IdealState idealState = getTableIdealState(tableNameWithType); + Preconditions.checkState(idealState != null, "Failed to find ideal state for table: %s", tableNameWithType); + Map> currentAssignment = idealState.getRecord().getMapFields(); + + return segmentAssignment.assignSegment(segmentName, currentAssignment, instancePartitionsMap); + } + private Map getInstacePartitionsMap(TableConfig tableConfig, @Nullable String tierName) { if (tierName != null) { diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/periodictask/SdsIntentCleanupTask.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/periodictask/SdsIntentCleanupTask.java new file mode 100644 index 0000000000..0448d82b5e --- /dev/null +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/periodictask/SdsIntentCleanupTask.java @@ -0,0 +1,103 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.controller.helix.core.periodictask; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import org.apache.helix.model.IdealState; +import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; +import org.apache.pinot.common.metrics.ControllerMetrics; +import org.apache.pinot.controller.LeadControllerManager; +import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; +import org.apache.pinot.spi.utils.CommonConstants; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/** + * Periodic task that cleans up orphaned SDS (Segment Distribution Service) upload intents. + * + *

Any segment with a non-null SDS status that hasn't reached UPLOAD_COMPLETE within 30 minutes + * is considered stale. If the segment is already in IdealState (partially completed), the status + * is fixed to UPLOAD_COMPLETE. Otherwise the ZK metadata is removed. + */ +public class SdsIntentCleanupTask extends ControllerPeriodicTask { + private static final Logger LOGGER = LoggerFactory.getLogger(SdsIntentCleanupTask.class); + + public static final String TASK_NAME = "SdsIntentCleanupTask"; + private static final long RUN_FREQUENCY_SECONDS = TimeUnit.MINUTES.toSeconds(10); + private static final long INITIAL_DELAY_SECONDS = TimeUnit.MINUTES.toSeconds(5); + private static final long STALE_THRESHOLD_MS = TimeUnit.MINUTES.toMillis(30); + + public SdsIntentCleanupTask(PinotHelixResourceManager pinotHelixResourceManager, + LeadControllerManager leadControllerManager, ControllerMetrics controllerMetrics) { + super(TASK_NAME, RUN_FREQUENCY_SECONDS, INITIAL_DELAY_SECONDS, pinotHelixResourceManager, + leadControllerManager, controllerMetrics); + } + + @Override + protected void processTable(String tableNameWithType) { + List segmentsZKMetadata = + _pinotHelixResourceManager.getSegmentsZKMetadata(tableNameWithType); + + long now = System.currentTimeMillis(); + IdealState idealState = null; + int cleaned = 0; + + for (SegmentZKMetadata segmentZKMetadata : segmentsZKMetadata) { + String sdsStatus = segmentZKMetadata.getSdsUploadStatus(); + if (sdsStatus == null || CommonConstants.Segment.SdsUploadStatus.UPLOAD_COMPLETE.equals(sdsStatus)) { + continue; + } + + long pushTime = segmentZKMetadata.getPushTime(); + if (pushTime <= 0 || (now - pushTime) <= STALE_THRESHOLD_MS) { + continue; + } + + String segmentName = segmentZKMetadata.getSegmentName(); + + // Lazy-load IdealState once per table + if (idealState == null) { + idealState = _pinotHelixResourceManager.getTableIdealState(tableNameWithType); + } + + // If segment is in IdealState, assignSegment() succeeded but status update failed. + // Fix the status instead of deleting — the segment may be actively serving. + Map instanceStateMap = + (idealState != null) ? idealState.getInstanceStateMap(segmentName) : null; + if (instanceStateMap != null && !instanceStateMap.isEmpty()) { + LOGGER.info("Fixing stuck SDS status for segment: {} in table: {} (was: {}, now: UPLOAD_COMPLETE)", + segmentName, tableNameWithType, sdsStatus); + segmentZKMetadata.setSdsUploadStatus(CommonConstants.Segment.SdsUploadStatus.UPLOAD_COMPLETE); + _pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata); + } else { + LOGGER.info("Removing orphaned SDS segment: {} in table: {} (status: {}, age: {}ms)", + segmentName, tableNameWithType, sdsStatus, now - pushTime); + _pinotHelixResourceManager.removeSegmentZKMetadata(tableNameWithType, segmentName); + } + cleaned++; + } + + if (cleaned > 0) { + LOGGER.info("Cleaned up {} orphaned SDS segments for table: {}", cleaned, tableNameWithType); + } + } +} diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerPeriodicTaskStarterStatelessTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerPeriodicTaskStarterStatelessTest.java index 413d23c383..c48e16a05f 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerPeriodicTaskStarterStatelessTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerPeriodicTaskStarterStatelessTest.java @@ -57,7 +57,7 @@ public ControllerStarter createControllerStarter() { } private class MockControllerStarter extends ControllerStarter { - private static final int NUM_PERIODIC_TASKS = 13; + private static final int NUM_PERIODIC_TASKS = 14; public MockControllerStarter() { super(); diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SdsTwoPhaseUploadIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SdsTwoPhaseUploadIntegrationTest.java new file mode 100644 index 0000000000..9b7f863a92 --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SdsTwoPhaseUploadIntegrationTest.java @@ -0,0 +1,392 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.integration.tests; + +import com.fasterxml.jackson.databind.JsonNode; +import java.io.File; +import java.io.OutputStream; +import java.net.HttpURLConnection; +import java.net.URL; +import java.nio.file.Files; +import java.util.List; +import java.util.Map; +import org.apache.commons.io.FileUtils; +import org.apache.helix.model.ExternalView; +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.controller.helix.ControllerTest; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.segment.spi.creator.SegmentIndexCreationDriver; +import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.JsonUtils; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.apache.pinot.spi.utils.builder.TableNameBuilder; +import org.apache.pinot.util.TestUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.fail; + + +/** + * Integration test for the SDS (Segment Distribution Service) 2-phase upload protocol. + * + *

Tests the full flow: + *

    + *
  1. Phase 1: uploadSegmentIntent (controller) — creates ZK metadata, computes assignment, returns upload URL
    2. + *
  2. Phase 2: sdsUpload (server) — uploads bytes to source server, server calls replicateSegment on controller
    3. + *
  3. Verification: segment is in IdealState, ExternalView shows ONLINE, queries return correct results
    4. + *
+ * + *

Also tests: + *

    + *
  • Conflict detection: uploading same segment again returns 409
    • + *
  • Traditional upload path still works alongside SDS (backward compatibility)
    • + *
+ */ +public class SdsTwoPhaseUploadIntegrationTest extends BaseClusterIntegrationTestSet { + private static final Logger LOGGER = LoggerFactory.getLogger(SdsTwoPhaseUploadIntegrationTest.class); + + private static final int NUM_BROKERS = 1; + private static final int NUM_SERVERS = 2; + private static final int REPLICATION_FACTOR = 2; + private static final String SDS_TEST_TABLE = "sdsTestTable"; + + private Schema _schema; + private TableConfig _tableConfig; + + @BeforeClass + public void setUp() + throws Exception { + TestUtils.ensureDirectoriesExistAndEmpty(_tempDir, _segmentDir, _tarDir); + + startZk(); + startController(); + startBrokers(NUM_BROKERS); + startServers(NUM_SERVERS); + + _schema = createSchema(); + addSchema(_schema); + + _tableConfig = createOfflineTableConfig(); + addTableConfig(_tableConfig); + + LOGGER.info("Setup complete. Cluster started with {} servers, replication={}", NUM_SERVERS, REPLICATION_FACTOR); + } + + @Override + protected String getTableName() { + return SDS_TEST_TABLE; + } + + @Override + protected TableConfig createOfflineTableConfig() { + return new TableConfigBuilder(TableType.OFFLINE) + .setTableName(getTableName()) + .setTimeColumnName(getTimeColumnName()) + .setNumReplicas(REPLICATION_FACTOR) + .build(); + } + + /** + * Test 1: Full SDS 2-phase upload happy path. + */ + @Test + public void testSdsTwoPhaseUploadHappyPath() + throws Exception { + LOGGER.info("Starting testSdsTwoPhaseUploadHappyPath"); + + String tableNameWithType = TableNameBuilder.OFFLINE.tableNameWithType(getTableName()); + + // Build a segment from avro + List avroFiles = unpackAvroData(_tempDir); + File avroFile = avroFiles.get(0); + + String segmentName = "sdsTestSegment_0"; + File segmentOutputDir = new File(_segmentDir, "sds_segment_0"); + SegmentGeneratorConfig segGenConfig = + new SegmentGeneratorConfig(_tableConfig, _schema); + segGenConfig.setInputFilePath(avroFile.getAbsolutePath()); + segGenConfig.setOutDir(segmentOutputDir.getAbsolutePath()); + segGenConfig.setSegmentName(segmentName); + segGenConfig.setTableName(tableNameWithType); + + SegmentIndexCreationDriver driver = new SegmentIndexCreationDriverImpl(); + driver.init(segGenConfig); + driver.build(); + + File segmentDir = new File(segmentOutputDir, segmentName); + assertTrue(segmentDir.exists(), "Segment directory should exist: " + segmentDir); + + // Read CRC for verification + SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(segmentDir); + long crc = Long.parseLong(segmentMetadata.getCrc()); + long totalDocs = segmentMetadata.getTotalDocs(); + + // Create a tar.gz of the segment + File tarFile = new File(_tarDir, segmentName + ".tar.gz"); + org.apache.pinot.common.utils.TarCompressionUtils.createCompressedTarFile(segmentDir, tarFile); + assertTrue(tarFile.exists(), "Tar file should exist"); + + // Build ZK metadata for the intent + // ZNRecord ID = segmentName; simple fields use CommonConstants.Segment keys + ZNRecord znRecord = new ZNRecord(segmentName); + znRecord.setSimpleField(CommonConstants.Segment.CRYPTER_NAME, ""); + znRecord.setSimpleField(CommonConstants.Segment.CRC, String.valueOf(crc)); + + // ===================== Phase 1: uploadSegmentIntent ===================== + LOGGER.info("Phase 1: Calling uploadSegmentIntent for segment: {}", segmentName); + + String intentUrl = _controllerBaseApiUrl + "/segments/" + getTableName() + + "/uploadIntent?tableType=OFFLINE"; + String intentResponse = ControllerTest.sendPostRequest(intentUrl, JsonUtils.objectToString(znRecord)); + JsonNode intentJson = JsonUtils.stringToJsonNode(intentResponse); + + LOGGER.info("Intent response: {}", intentResponse); + + assertEquals(intentJson.get("segmentName").asText(), segmentName); + assertEquals(intentJson.get("tableNameWithType").asText(), tableNameWithType); + assertNotNull(intentJson.get("sourceServer").asText()); + assertTrue(intentJson.get("assignedInstances").size() > 0, "Should have assigned instances"); + String sourceServerUploadUrl = intentJson.get("sourceServerUploadUrl").asText(); + assertNotNull(sourceServerUploadUrl, "Should have upload URL"); + assertTrue(sourceServerUploadUrl.contains("/sdsUpload"), "URL should contain sdsUpload path"); + + LOGGER.info("Got sourceServerUploadUrl: {}", sourceServerUploadUrl); + + // Verify ZK metadata was created with INTENT_CREATED status + SegmentZKMetadata zkMetadata = _helixResourceManager.getSegmentZKMetadata(tableNameWithType, segmentName); + assertNotNull(zkMetadata, "ZK metadata should exist after intent"); + assertEquals(zkMetadata.getSdsUploadStatus(), CommonConstants.Segment.SdsUploadStatus.INTENT_CREATED); + assertEquals(zkMetadata.getDownloadUrl(), CommonConstants.Segment.METADATA_URI_FOR_PEER_DOWNLOAD); + + // Verify segment is NOT yet in IdealState + IdealState idealState = _helixResourceManager.getTableIdealState(tableNameWithType); + assertNotNull(idealState); + assertNull(idealState.getInstanceStateMap(segmentName), + "Segment should NOT be in IdealState before Phase 2"); + + // ===================== Phase 2: sdsUpload to source server ===================== + LOGGER.info("Phase 2: Uploading segment bytes to source server at: {}", sourceServerUploadUrl); + + String uploadUrl = sourceServerUploadUrl + "?expectedCrc=" + crc + "&minReplicas=1&timeoutMs=120000"; + HttpURLConnection conn = (HttpURLConnection) new URL(uploadUrl).openConnection(); + conn.setRequestMethod("POST"); + conn.setDoOutput(true); + conn.setRequestProperty("Content-Type", "application/octet-stream"); + conn.setConnectTimeout(30000); + conn.setReadTimeout(180000); + + byte[] segmentBytes = Files.readAllBytes(tarFile.toPath()); + try (OutputStream os = conn.getOutputStream()) { + os.write(segmentBytes); + } + + int statusCode = conn.getResponseCode(); + String responseBody = ""; + try { + responseBody = new String(conn.getInputStream().readAllBytes()); + } catch (Exception e) { + if (conn.getErrorStream() != null) { + responseBody = new String(conn.getErrorStream().readAllBytes()); + } + } + conn.disconnect(); + + LOGGER.info("SDS upload response: status={}, body={}", statusCode, responseBody); + assertEquals(statusCode, 200, "SDS upload should succeed. Response: " + responseBody); + + // ===================== Verification ===================== + LOGGER.info("Verifying segment is in IdealState and ExternalView"); + + // Verify segment is now in IdealState with REPLICATION_FACTOR instances + TestUtils.waitForCondition(aVoid -> { + IdealState is = _helixResourceManager.getTableIdealState(tableNameWithType); + if (is == null) { + return false; + } + Map stateMap = is.getInstanceStateMap(segmentName); + return stateMap != null && stateMap.size() == REPLICATION_FACTOR; + }, 30_000L, "Segment should be in IdealState with " + REPLICATION_FACTOR + " replicas"); + + // Verify ZK metadata status is UPLOAD_COMPLETE + zkMetadata = _helixResourceManager.getSegmentZKMetadata(tableNameWithType, segmentName); + assertEquals(zkMetadata.getSdsUploadStatus(), CommonConstants.Segment.SdsUploadStatus.UPLOAD_COMPLETE); + assertNotNull(zkMetadata.getSourceServer(), "sourceServer should be set in ZK metadata"); + + // Verify all replicas come ONLINE in ExternalView + TestUtils.waitForCondition(aVoid -> { + ExternalView ev = _helixAdmin.getResourceExternalView(getHelixClusterName(), tableNameWithType); + if (ev == null) { + return false; + } + Map stateMap = ev.getStateMap(segmentName); + if (stateMap == null) { + return false; + } + long onlineCount = stateMap.values().stream().filter("ONLINE"::equals).count(); + LOGGER.info("Segment {} has {} ONLINE replicas", segmentName, onlineCount); + return onlineCount >= REPLICATION_FACTOR; + }, 120_000L, "All replicas should come ONLINE in ExternalView"); + + // Verify queries work + TestUtils.waitForCondition(aVoid -> { + try { + long count = getCurrentCountStarResult(); + LOGGER.info("COUNT(*) = {} (expected {})", count, totalDocs); + return count == totalDocs; + } catch (Exception e) { + return false; + } + }, 30_000L, "Query should return correct document count"); + + LOGGER.info("testSdsTwoPhaseUploadHappyPath completed successfully"); + } + + /** + * Test 2: Uploading the same segment again should return 409 CONFLICT. + */ + @Test(dependsOnMethods = "testSdsTwoPhaseUploadHappyPath") + public void testSdsUploadConflict() + throws Exception { + LOGGER.info("Starting testSdsUploadConflict"); + + String segmentName = "sdsTestSegment_0"; + String tableNameWithType = TableNameBuilder.OFFLINE.tableNameWithType(getTableName()); + + // Build ZK metadata for the same segment (ID = segmentName) + ZNRecord znRecord = new ZNRecord(segmentName); + + String intentUrl = _controllerBaseApiUrl + "/segments/" + getTableName() + + "/uploadIntent?tableType=OFFLINE"; + + try { + ControllerTest.sendPostRequest(intentUrl, JsonUtils.objectToString(znRecord)); + fail("Should have thrown exception for duplicate segment"); + } catch (Exception e) { + LOGGER.info("Got expected error for duplicate segment: {}", e.getMessage()); + assertTrue(e.getMessage().contains("409") || e.getMessage().contains("already exists"), + "Should get 409 CONFLICT, got: " + e.getMessage()); + } + + LOGGER.info("testSdsUploadConflict completed successfully"); + } + + /** + * Test 3: Traditional upload still works alongside SDS (backward compatibility). + * Upload a different segment via the standard upload API and verify both segments are queryable. + */ + @Test(dependsOnMethods = "testSdsUploadConflict") + public void testTraditionalUploadStillWorks() + throws Exception { + LOGGER.info("Starting testTraditionalUploadStillWorks"); + + String tableNameWithType = TableNameBuilder.OFFLINE.tableNameWithType(getTableName()); + + // Build a second segment + List avroFiles = getAllAvroFiles(); + File avroFile = avroFiles.get(1); + + String segmentName = "traditionalSegment_0"; + File segmentOutputDir = new File(_segmentDir, "traditional_segment_0"); + SegmentGeneratorConfig segGenConfig = + new SegmentGeneratorConfig(_tableConfig, _schema); + segGenConfig.setInputFilePath(avroFile.getAbsolutePath()); + segGenConfig.setOutDir(segmentOutputDir.getAbsolutePath()); + segGenConfig.setSegmentName(segmentName); + segGenConfig.setTableName(tableNameWithType); + + SegmentIndexCreationDriver driver = new SegmentIndexCreationDriverImpl(); + driver.init(segGenConfig); + driver.build(); + + File segmentDir = new File(segmentOutputDir, segmentName); + SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(segmentDir); + long newDocs = segmentMetadata.getTotalDocs(); + + // Create tar and upload via traditional path + File tarFile = new File(_tarDir, segmentName + ".tar.gz"); + org.apache.pinot.common.utils.TarCompressionUtils.createCompressedTarFile(segmentDir, tarFile); + uploadSegments(getTableName(), _tarDir); + + // Wait for both segments to be loaded + TestUtils.waitForCondition(aVoid -> { + IdealState is = _helixResourceManager.getTableIdealState(tableNameWithType); + return is != null && is.getRecord().getMapFields().size() == 2; + }, 30_000L, "Should have 2 segments in IdealState"); + + // Verify the SDS segment still has SDS metadata + SegmentZKMetadata sdsMetadata = + _helixResourceManager.getSegmentZKMetadata(tableNameWithType, "sdsTestSegment_0"); + assertNotNull(sdsMetadata.getSdsUploadStatus(), "SDS segment should retain SDS status"); + assertEquals(sdsMetadata.getSdsUploadStatus(), CommonConstants.Segment.SdsUploadStatus.UPLOAD_COMPLETE); + + // Verify the traditional segment does NOT have SDS metadata + SegmentZKMetadata traditionalMetadata = + _helixResourceManager.getSegmentZKMetadata(tableNameWithType, segmentName); + assertNull(traditionalMetadata.getSdsUploadStatus(), + "Traditional segment should NOT have SDS status"); + + // Verify total query count includes both segments + // Get the SDS segment's doc count + SegmentZKMetadata sdsZk = _helixResourceManager.getSegmentZKMetadata(tableNameWithType, "sdsTestSegment_0"); + long sdsSegmentDocs = sdsZk.getTotalDocs(); + + TestUtils.waitForCondition(aVoid -> { + try { + long count = getCurrentCountStarResult(); + long expected = sdsSegmentDocs + newDocs; + LOGGER.info("COUNT(*) = {} (expected {})", count, expected); + return count == expected; + } catch (Exception e) { + return false; + } + }, 60_000L, "Query should return combined document count from both segments"); + + LOGGER.info("testTraditionalUploadStillWorks completed successfully"); + } + + @AfterClass + public void tearDown() + throws Exception { + try { + dropOfflineTable(SDS_TEST_TABLE); + } catch (Exception e) { + LOGGER.warn("Failed to drop table: {}", e.getMessage()); + } + stopServer(); + stopBroker(); + stopController(); + stopZk(); + FileUtils.deleteDirectory(_tempDir); + } +} diff --git a/pinot-server/src/main/java/org/apache/pinot/server/api/resources/TablesResource.java b/pinot-server/src/main/java/org/apache/pinot/server/api/resources/TablesResource.java index da822cef86..363d220f8d 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/api/resources/TablesResource.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/api/resources/TablesResource.java @@ -30,7 +30,11 @@ import io.swagger.annotations.SwaggerDefinition; import java.io.File; import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.net.HttpURLConnection; import java.net.URI; +import java.net.URL; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.util.ArrayList; @@ -43,6 +47,7 @@ import java.util.stream.Collectors; import javax.inject.Inject; import javax.inject.Named; +import javax.ws.rs.Consumes; import javax.ws.rs.DefaultValue; import javax.ws.rs.Encoded; import javax.ws.rs.GET; @@ -58,7 +63,11 @@ import javax.ws.rs.core.Response; import javax.ws.rs.core.StreamingOutput; import org.apache.commons.io.FileUtils; +import org.apache.commons.io.IOUtils; import org.apache.commons.lang3.tuple.Pair; +import org.apache.helix.HelixAdmin; +import org.apache.helix.HelixManager; +import org.apache.helix.model.ExternalView; import org.apache.helix.model.IdealState; import org.apache.pinot.common.metadata.ZKMetadataProvider; import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; @@ -98,6 +107,7 @@ import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; import org.apache.pinot.server.access.AccessControlFactory; import org.apache.pinot.server.api.AdminApiApplication; +import org.apache.pinot.server.realtime.ControllerLeaderLocator; import org.apache.pinot.server.starter.ServerInstance; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.data.FieldSpec; @@ -459,7 +469,6 @@ public Response downloadSegment( Response.Status.NOT_FOUND); } try { - // TODO Limit the number of concurrent downloads of segments because compression is an expensive operation. // Store the tar.gz segment file in the server's segmentTarDir folder with a unique file name. // Note that two clients asking the same segment file will result in the same tar.gz files being created twice. // Will revisit for optimization if performance becomes an issue. @@ -489,6 +498,207 @@ public Response downloadSegment( } } + @POST + @Produces(MediaType.APPLICATION_JSON) + @Consumes(MediaType.APPLICATION_OCTET_STREAM) + @Path("/segments/{tableNameWithType}/{segmentName}/sdsUpload") + @ApiOperation(value = "Upload a segment directly to this server via SDS", + notes = "SDS Phase 2: Receives segment bytes, stores locally, then calls controller to replicate.") + @ApiResponses(value = { + @ApiResponse(code = 200, message = "Successfully uploaded and replicated segment"), + @ApiResponse(code = 400, message = "Bad Request"), + @ApiResponse(code = 500, message = "Internal error"), + @ApiResponse(code = 504, message = "Timeout waiting for minimum replicas") + }) + public Response uploadSegmentToServer( + @ApiParam(value = "Name of the table with type", required = true) + @PathParam("tableNameWithType") String tableNameWithType, + @ApiParam(value = "Name of the segment", required = true) + @PathParam("segmentName") @Encoded String segmentName, + @ApiParam(value = "Expected CRC of the segment") + @QueryParam("expectedCrc") @DefaultValue("-1") long expectedCrc, + @ApiParam(value = "Minimum number of replicas to wait for before ACK") + @QueryParam("minReplicas") @DefaultValue("1") int minReplicas, + @ApiParam(value = "Timeout in milliseconds for waiting for replicas") + @QueryParam("timeoutMs") @DefaultValue("300000") long timeoutMs, + InputStream segmentInputStream, + @Context HttpHeaders httpHeaders) + throws Exception { + + tableNameWithType = DatabaseUtils.translateTableName(tableNameWithType, httpHeaders); + segmentName = URIUtils.decode(segmentName); + LOGGER.info("Received SDS upload request for segment: {} of table: {}", segmentName, tableNameWithType); + + // Validate access + ServerResourceUtils.validateDataAccess(_accessControlFactory, tableNameWithType, httpHeaders); + + // Get data directory from instance data manager. + // Note: we intentionally do NOT require TableDataManager to exist here because in the SDS flow + // the source server receives segment bytes BEFORE it appears in IdealState, so the + // TableDataManager may not have been created yet. + InstanceDataManager instanceDataManager = _serverInstance.getInstanceDataManager(); + String dataDir = instanceDataManager.getInstanceDataDir(); + + // Write to a staging directory first, then atomically move to the final data directory. + // This ensures tryLoadExistingSegment() only sees complete segments on disk. + // Staging dirs are cleaned up on server restart by HelixInstanceDataManager.cleanupStagingDirectories(). + File tableDir = new File(dataDir, tableNameWithType); + File stagingDir = new File(tableDir, "staging"); + stagingDir.mkdirs(); + File segmentStagingDir = new File(stagingDir, segmentName + "_" + UUID.randomUUID()); + segmentStagingDir.mkdirs(); + + try { + // Write segment tar to staging + File segmentTarFile = new File(segmentStagingDir, + segmentName + TarCompressionUtils.TAR_COMPRESSED_FILE_EXTENSION); + LOGGER.info("Writing segment: {} to staging: {}", segmentName, segmentTarFile); + try (OutputStream os = Files.newOutputStream(segmentTarFile.toPath())) { + IOUtils.copy(segmentInputStream, os); + } + LOGGER.info("Wrote segment tar: {} ({} bytes)", segmentTarFile, segmentTarFile.length()); + + // Untar segment + File untarDir = new File(segmentStagingDir, "untarred"); + File untarredSegmentDir = TarCompressionUtils.untar(segmentTarFile, untarDir).get(0); + LOGGER.info("Untarred segment: {} to: {}", segmentName, untarredSegmentDir); + + // Verify CRC if provided + if (expectedCrc >= 0) { + SegmentMetadataImpl metadata = new SegmentMetadataImpl(untarredSegmentDir); + long actualCrc = Long.parseLong(metadata.getCrc()); + if (expectedCrc != actualCrc) { + throw new WebApplicationException( + String.format("CRC mismatch for segment: %s, expected: %d, actual: %d", + segmentName, expectedCrc, actualCrc), + Response.Status.BAD_REQUEST); + } + LOGGER.info("CRC verification passed for segment: {}", segmentName); + } + + // Move segment to final data directory BEFORE calling replicateSegment(). + // This way, when the controller updates IdealState and this server gets OFFLINE→ONLINE, + // tryLoadExistingSegment() will find the segment on disk and load it without downloading. + File finalSegmentDir = new File(tableDir, segmentName); + if (finalSegmentDir.exists()) { + FileUtils.deleteDirectory(finalSegmentDir); + } + FileUtils.moveDirectory(untarredSegmentDir, finalSegmentDir); + LOGGER.info("Moved segment: {} to final location: {}", segmentName, finalSegmentDir); + + // Notify the controller to update IdealState, triggering OFFLINE→ONLINE on all replicas + callControllerReplicateSegment(tableNameWithType, segmentName); + + // Block until enough replicas are ONLINE in ExternalView before ACKing the client + waitForMinimumReplicasOnline(tableNameWithType, segmentName, minReplicas, timeoutMs); + + LOGGER.info("SDS upload completed for segment: {} of table: {}", segmentName, tableNameWithType); + String successMsg = "Successfully uploaded segment: " + segmentName + " to table: " + tableNameWithType; + return Response.ok(JsonUtils.objectToString(Map.of("status", successMsg))).build(); + } catch (WebApplicationException e) { + throw e; + } catch (Exception e) { + LOGGER.error("Failed to upload segment: {} for table: {}", segmentName, tableNameWithType, e); + throw new WebApplicationException( + "Failed to upload segment: " + segmentName + ": " + e.getMessage(), + Response.Status.INTERNAL_SERVER_ERROR); + } finally { + // Clean up staging directory + FileUtils.deleteQuietly(segmentStagingDir); + } + } + + /** + * Calls the controller's replicateSegment API to notify that a segment has been uploaded to this server. + */ + private void callControllerReplicateSegment(String tableNameWithType, String segmentName) + throws Exception { + // Get controller leader URL + String rawTableName = TableNameBuilder.extractRawTableName(tableNameWithType); + ControllerLeaderLocator locator = ControllerLeaderLocator.getInstance(); + Pair controllerLeader = locator.getControllerLeader(rawTableName); + if (controllerLeader == null) { + throw new WebApplicationException("Cannot find controller leader for table: " + tableNameWithType, + Response.Status.INTERNAL_SERVER_ERROR); + } + + String tableType = TableNameBuilder.getTableTypeFromTableName(tableNameWithType) == TableType.OFFLINE + ? "OFFLINE" : "REALTIME"; + + // Build the replicateSegment URL + String controllerUrl = "http://" + controllerLeader.getLeft() + ":" + controllerLeader.getRight() + + "/segments/" + rawTableName + "/" + URIUtils.encode(segmentName) + + "/replicate?tableType=" + tableType; + + LOGGER.info("Calling controller replicateSegment at: {}", controllerUrl); + + // Make HTTP POST call using java.net.HttpURLConnection + HttpURLConnection connection = (HttpURLConnection) new URL(controllerUrl).openConnection(); + try { + connection.setRequestMethod("POST"); + connection.setConnectTimeout(30000); + connection.setReadTimeout(60000); + + int statusCode = connection.getResponseCode(); + if (statusCode != HttpURLConnection.HTTP_OK) { + String responseBody = ""; + try (InputStream errorStream = connection.getErrorStream()) { + if (errorStream != null) { + responseBody = new String(errorStream.readAllBytes(), StandardCharsets.UTF_8); + } + } + throw new WebApplicationException( + "Controller replicateSegment failed with status: " + statusCode + ", body: " + responseBody, + Response.Status.INTERNAL_SERVER_ERROR); + } + LOGGER.info("Controller replicateSegment succeeded for segment: {} of table: {}", + segmentName, tableNameWithType); + } finally { + connection.disconnect(); + } + } + + /** + * Waits for the minimum number of replicas of a segment to come ONLINE in the ExternalView. + */ + private void waitForMinimumReplicasOnline(String tableNameWithType, String segmentName, + int minReplicas, long timeoutMs) + throws Exception { + LOGGER.info("Waiting for {} replicas of segment: {} to come ONLINE (timeout: {}ms)", + minReplicas, segmentName, timeoutMs); + + HelixManager helixManager = _serverInstance.getHelixManager(); + HelixAdmin helixAdmin = helixManager.getClusterManagmentTool(); + String clusterName = helixManager.getClusterName(); + + long deadline = System.currentTimeMillis() + timeoutMs; + int pollIntervalMs = 2000; + + while (System.currentTimeMillis() < deadline) { + ExternalView externalView = helixAdmin.getResourceExternalView(clusterName, tableNameWithType); + if (externalView != null) { + Map stateMap = externalView.getStateMap(segmentName); + if (stateMap != null) { + long onlineCount = stateMap.values().stream() + .filter(SegmentStateModel.ONLINE::equals) + .count(); + if (onlineCount >= minReplicas) { + LOGGER.info("Segment: {} has {} ONLINE replicas (required: {})", + segmentName, onlineCount, minReplicas); + return; + } + LOGGER.info("Segment: {} has {} ONLINE replicas, waiting for {} (deadline in {}ms)", + segmentName, onlineCount, minReplicas, deadline - System.currentTimeMillis()); + } + } + Thread.sleep(pollIntervalMs); + } + + throw new WebApplicationException( + String.format("Timeout waiting for %d replicas of segment: %s to come ONLINE", minReplicas, segmentName), + Response.Status.GATEWAY_TIMEOUT); + } + @GET @Produces(MediaType.APPLICATION_JSON) @Path("/segments/{tableNameWithType}/{segmentName}/validDocIdsBitmap") diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManager.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManager.java index 6785d3db58..e99c79596e 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManager.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManager.java @@ -235,6 +235,9 @@ public String getInstanceId() { @Override public synchronized void start() { + // Clean up staging directories that may have been left behind by interrupted SDS uploads + cleanupStagingDirectories(); + _propertyStore = _helixManager.getHelixPropertyStore(); // Initialize logical table metadata cache _logicalTableMetadataCache.init(_propertyStore); @@ -242,6 +245,32 @@ public synchronized void start() { LOGGER.info("Helix instance data manager started"); } + /** + * Cleans up staging directories left behind by interrupted SDS uploads. + * Pattern: {dataDir}/{tableNameWithType}/staging/ + */ + private void cleanupStagingDirectories() { + try { + File instanceDataDir = new File(_instanceDataManagerConfig.getInstanceDataDir()); + if (!instanceDataDir.exists()) { + return; + } + File[] tableDataDirs = instanceDataDir.listFiles((dir, name) -> TableNameBuilder.isTableResource(name)); + if (tableDataDirs == null) { + return; + } + for (File tableDataDir : tableDataDirs) { + File stagingDir = new File(tableDataDir, "staging"); + if (stagingDir.exists()) { + LOGGER.info("Deleting staging directory: {}", stagingDir); + FileUtils.deleteQuietly(stagingDir); + } + } + } catch (Exception e) { + LOGGER.error("Failed to clean up staging directories", e); + } + } + @Override public synchronized void shutDown() { _segmentReloadExecutor.shutdownNow(); diff --git a/pinot-server/src/test/java/org/apache/pinot/server/api/BaseResourceTest.java b/pinot-server/src/test/java/org/apache/pinot/server/api/BaseResourceTest.java index d86a02f0ac..0c3d7e6b40 100644 --- a/pinot-server/src/test/java/org/apache/pinot/server/api/BaseResourceTest.java +++ b/pinot-server/src/test/java/org/apache/pinot/server/api/BaseResourceTest.java @@ -118,6 +118,7 @@ public void setUp() when(serverInstance.getInstanceDataManager()).thenReturn(instanceDataManager); when(serverInstance.getInstanceDataManager().getSegmentFileDirectory()).thenReturn( FileUtils.getTempDirectoryPath()); + when(instanceDataManager.getInstanceDataDir()).thenReturn(TEMP_DIR.getAbsolutePath()); when(serverInstance.getHelixManager()).thenReturn(mock(HelixManager.class)); // Mock the segment uploader diff --git a/pinot-server/src/test/java/org/apache/pinot/server/api/TablesResourceTest.java b/pinot-server/src/test/java/org/apache/pinot/server/api/TablesResourceTest.java index 0d208cc695..6815c1612d 100644 --- a/pinot-server/src/test/java/org/apache/pinot/server/api/TablesResourceTest.java +++ b/pinot-server/src/test/java/org/apache/pinot/server/api/TablesResourceTest.java @@ -26,6 +26,7 @@ import java.util.List; import java.util.Map; import javax.ws.rs.client.Entity; +import javax.ws.rs.core.MediaType; import javax.ws.rs.core.Response; import org.apache.commons.configuration2.ex.ConfigurationException; import org.apache.commons.io.FileUtils; @@ -592,4 +593,96 @@ public void testOfflineTableSegmentMetadata() .request().get(Response.class); Assert.assertEquals(response.getStatus(), Response.Status.NOT_FOUND.getStatusCode()); } + + // ==================== SDS Upload Tests ==================== + + /** + * Helper to create a segment tar file from an existing segment's index directory. + */ + private File createSegmentTarFile(ImmutableSegment segment) + throws IOException { + File segmentIndexDir = segment.getSegmentMetadata().getIndexDir(); + File tarFile = new File(FileUtils.getTempDirectory(), + segment.getSegmentName() + TarCompressionUtils.TAR_GZ_FILE_EXTENSION); + TarCompressionUtils.createCompressedTarFile(segmentIndexDir, tarFile); + return tarFile; + } + + @Test + public void testSdsUploadNonExistentTable() + throws Exception { + // In SDS flow, the source server receives segment bytes BEFORE it appears in IdealState, + // so the TableDataManager may not exist yet. The endpoint no longer returns 404 for a + // non-existent table; instead it proceeds and fails later at the controller replicateSegment call. + ImmutableSegment segment = _offlineIndexSegments.get(0); + File tarFile = createSegmentTarFile(segment); + try { + byte[] segmentBytes = FileUtils.readFileToByteArray(tarFile); + + // POST to a table that does not exist on this server — should fail with 500 + // because the controller replicateSegment call will fail (no controller in unit test) + Response response = _webTarget.path( + String.format("/segments/%s/%s/sdsUpload", "nonExistentTable_OFFLINE", segment.getSegmentName())) + .request() + .post(Entity.entity(segmentBytes, MediaType.APPLICATION_OCTET_STREAM_TYPE), Response.class); + + Assert.assertEquals(response.getStatus(), Response.Status.INTERNAL_SERVER_ERROR.getStatusCode()); + } finally { + FileUtils.deleteQuietly(tarFile); + } + } + + @Test + public void testSdsUploadCrcMismatch() + throws Exception { + // Use an existing offline segment to create a tar file + ImmutableSegment segment = _offlineIndexSegments.get(0); + File tarFile = createSegmentTarFile(segment); + try { + byte[] segmentBytes = FileUtils.readFileToByteArray(tarFile); + String actualCrc = segment.getSegmentMetadata().getCrc(); + // Use a CRC value that definitely does not match the segment's actual CRC + long wrongCrc = Long.parseLong(actualCrc) + 1; + + String offlineTableName = TableNameBuilder.OFFLINE.tableNameWithType(TABLE_NAME); + Response response = _webTarget.path( + String.format("/segments/%s/%s/sdsUpload", offlineTableName, segment.getSegmentName())) + .queryParam("expectedCrc", wrongCrc) + .request() + .post(Entity.entity(segmentBytes, MediaType.APPLICATION_OCTET_STREAM_TYPE), Response.class); + + Assert.assertEquals(response.getStatus(), Response.Status.BAD_REQUEST.getStatusCode()); + } finally { + FileUtils.deleteQuietly(tarFile); + } + } + + @Test + public void testSdsUploadSuccessUntilControllerCall() + throws Exception { + // This test validates that the SDS upload endpoint correctly: + // 1. Receives and writes the segment tar to staging + // 2. Untars the segment + // 3. Moves the segment to the final data directory + // The test expects a 500 because there is no real controller to handle the replicateSegment call. + ImmutableSegment segment = _offlineIndexSegments.get(0); + File tarFile = createSegmentTarFile(segment); + try { + byte[] segmentBytes = FileUtils.readFileToByteArray(tarFile); + String offlineTableName = TableNameBuilder.OFFLINE.tableNameWithType(TABLE_NAME); + + // Do not pass expectedCrc so CRC verification is skipped (default is -1) + Response response = _webTarget.path( + String.format("/segments/%s/%s/sdsUpload", offlineTableName, segment.getSegmentName())) + .request() + .post(Entity.entity(segmentBytes, MediaType.APPLICATION_OCTET_STREAM_TYPE), Response.class); + + // Expect 500 because ControllerLeaderLocator is not set up in the test environment, + // so callControllerReplicateSegment() will fail. This validates that the upload, untar, + // and move-to-final-directory steps all completed successfully before reaching the controller call. + Assert.assertEquals(response.getStatus(), Response.Status.INTERNAL_SERVER_ERROR.getStatusCode()); + } finally { + FileUtils.deleteQuietly(tarFile); + } + } } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index bf9a8f0f57..48748b7382 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -1544,6 +1544,15 @@ public static class Offline { public static final String CUSTOM_MAP = "custom.map"; public static final String SIZE_IN_BYTES = "segment.size.in.bytes"; public static final String SOURCE_SERVER = "segment.source.server"; + // SDS (Segment Distribution Service) 2-phase upload status, tracked in ZK metadata. + // Flow: INTENT_CREATED (Phase 1) → UPLOAD_PROCESSING (replicating) → UPLOAD_COMPLETE + public static final String SDS_UPLOAD_STATUS = "segment.sds.upload.status"; + + public static class SdsUploadStatus { + public static final String INTENT_CREATED = "INTENT_CREATED"; + public static final String UPLOAD_PROCESSING = "UPLOAD_PROCESSING"; + public static final String UPLOAD_COMPLETE = "UPLOAD_COMPLETE"; + } /** * This field is used for parallel push protection to lock the segment globally.