Implement registerSegment API to support corp to prod migration

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)

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",

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