Implement SDS two-phase upload and segment distribution flow

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);
   }

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<PeriodicTask> 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;
   }

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<String, List<FormDataBodyPart>> 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<String> 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);
+    }
+  }
 }

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<String> _assignedInstances;
+  private final String _sourceServer;
+  private final String _sourceServerUploadUrl;
+
+  @JsonCreator
+  public UploadSegmentIntentResponse(
+      @JsonProperty("segmentName") String segmentName,
+      @JsonProperty("tableNameWithType") String tableNameWithType,
+      @JsonProperty("assignedInstances") List<String> 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<String> getAssignedInstances() {
+    return _assignedInstances;
+  }
+
+  @JsonProperty("sourceServer")
+  public String getSourceServer() {
+    return _sourceServer;
+  }
+
+  @JsonProperty("sourceServerUploadUrl")
+  public String getSourceServerUploadUrl() {
+    return _sourceServerUploadUrl;
+  }
+}