Add a register endpoint to upload zk metadata directly by dinoocch · Pull Request #145 · linkedin/pinot

dinoocch · 2026-03-12T18:31:16Z

Background

For corp to prod we need to move all the corp data and split it up into several prod helix clusters.
They read/write from separate gpfs paths.

To help speed up the migration we pre-copied the data into the destination and are running a continuous rsync.

Our primary limiting factors today will be:

GPFS Load
Controller Load

GPFS load is especially dangerous since the underlying gpfs cluster is shared with our prod-lva1 service unit.

We have previously discussed the "peer-download" to reduce the server side download reads to 1. This proposal focuses on the upload path (before servers start to download).

For segment upload we have lots of options today:

Multipart using the built segment tar.gz
- Issue: Can't be used since we have already copied the data + will unpack on the controller to calculate metadata
URI based upload where the destination controller can download from a URI / Peer Controller
- Issue: This doesn't help us with gpfs reads
File URI -- similar to the peer upload
Metadata upload (there the segment data and metadata are uploaded separately)

Idea: Upload the segment zk metadata and assign the segment

Copilot

Pull request overview

Adds a new controller REST endpoint to register (or refresh) a segment directly from pre-built ZK metadata JSON, avoiding segment file upload/unpacking/copying to reduce controller/GPFS load during migration workflows.

Changes:

Added POST /segments/register that accepts SegmentZKMetadata JSON and registers/refreshes the segment in ZK.
Implemented logic to recover from partial prior registrations by removing “stale” segment ZK metadata when IdealState is missing the segment.
On new registration: optionally updates target tier, writes segment ZK metadata, and assigns the segment into Helix IdealState.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

Copilot · 2026-03-12T18:35:07Z

+      segmentZKMetadata.setRefreshTime(System.currentTimeMillis());
+      if (!_pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata,


On the refresh path, the code updates ZK metadata using the request-provided SegmentZKMetadata object. If the request body omits existing simpleFields/customMap, this will overwrite the existing ZNRecord and effectively delete previously stored metadata fields. Consider loading the existing SegmentZKMetadata from existingZNRecord and merging/whitelisting only the fields that are allowed to change (e.g., downloadUrl, crc, refreshTime), rather than replacing the whole record.

Suggested change

segmentZKMetadata.setRefreshTime(System.currentTimeMillis());

if (!_pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata,

// Merge request-provided fields into existing ZK metadata instead of overwriting the whole record.

SegmentZKMetadata mergedSegmentZKMetadata = new SegmentZKMetadata(existingZNRecord);

// Whitelist fields that are allowed to change on refresh.

mergedSegmentZKMetadata.setDownloadUrl(segmentZKMetadata.getDownloadUrl());

mergedSegmentZKMetadata.setCrc(segmentZKMetadata.getCrc());

mergedSegmentZKMetadata.setRefreshTime(System.currentTimeMillis());

if (!_pinotHelixResourceManager.updateZkMetadata(tableNameWithType, mergedSegmentZKMetadata,

Copilot · 2026-03-12T18:35:08Z

+    if (StringUtils.isEmpty(segmentZKMetadata.getDownloadUrl())) {
+      throw new ControllerApplicationException(LOGGER,
+          "segment.download.url is required in simpleFields", Response.Status.BAD_REQUEST);
+    }


The swagger notes state that segment.crc is required, but the implementation only validates segment.download.url. If crc is missing, SegmentZKMetadata#getCrc() will be -1 and downstream refresh/download logic can behave incorrectly. Please validate that segmentZKMetadata.getCrc() is present/valid (>= 0) and return 400 if not.

Suggested change

}

}

long crc = segmentZKMetadata.getCrc();

if (crc < 0) {

throw new ControllerApplicationException(LOGGER,

"segment.crc is required in simpleFields and must be non-negative", Response.Status.BAD_REQUEST);

}

Copilot · 2026-03-12T18:35:08Z

+  private SuccessResponse registerSegmentInternal(String segmentZKMetadataJson, String tableName, String tableType,
+      boolean allowRefresh) {
+    if (StringUtils.isEmpty(tableName)) {
+      throw new ControllerApplicationException(LOGGER, "tableName is required", Response.Status.BAD_REQUEST);
+    }
+
+    SegmentZKMetadata segmentZKMetadata;
+    try {
+      segmentZKMetadata = SegmentZKMetadata.fromJsonString(segmentZKMetadataJson);
+    } catch (Exception e) {
+      throw new ControllerApplicationException(LOGGER,
+          "Failed to parse segment ZK metadata: " + e.getMessage(), Response.Status.BAD_REQUEST);
+    }
+
+    String segmentName = segmentZKMetadata.getSegmentName();
+    if (StringUtils.isEmpty(segmentName)) {
+      throw new ControllerApplicationException(LOGGER, "segmentName is required in ZK metadata",
+          Response.Status.BAD_REQUEST);
+    }
+    if (StringUtils.isEmpty(segmentZKMetadata.getDownloadUrl())) {
+      throw new ControllerApplicationException(LOGGER,
+          "segment.download.url is required in simpleFields", Response.Status.BAD_REQUEST);
+    }
+
+    String rawTableName = TableNameBuilder.extractRawTableName(tableName);
+    TableType type = TableType.valueOf(tableType.toUpperCase());
+    String tableNameWithType = TableNameBuilder.forType(type).tableNameWithType(rawTableName);
+
+    TableConfig tableConfig = _pinotHelixResourceManager.getTableConfig(tableNameWithType);
+    if (tableConfig == null) {
+      throw new ControllerApplicationException(LOGGER, "Failed to find table: " + tableNameWithType,
+          Response.Status.BAD_REQUEST);
+    }
+
+    LOGGER.info("Registering segment: {} for table: {}", segmentName, tableNameWithType);
+
+    ZNRecord existingZNRecord = _pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, segmentName);
+
+    // If ZK metadata exists but the segment is absent from the ideal state, a previous registration failed
+    // partway through. Remove the stale ZK entry and treat this as a new segment.
+    if (existingZNRecord != null) {
+      IdealState idealState = _pinotHelixResourceManager.getTableIdealState(tableNameWithType);
+      Preconditions.checkState(idealState != null, "Failed to find ideal state for table: %s", tableNameWithType);
+      if (idealState.getInstanceStateMap(segmentName) == null) {
+        LOGGER.warn("Removing stale segment ZK metadata (recovering from previous registration failure) "
+            + "for table: {}, segment: {}", tableNameWithType, segmentName);
+        Preconditions.checkState(_pinotHelixResourceManager.removeSegmentZKMetadata(tableNameWithType, segmentName),
+            "Failed to remove stale segment ZK metadata for table: %s, segment: %s", tableNameWithType, segmentName);
+        existingZNRecord = null;
+      }
+    }
+
+    if (existingZNRecord != null) {
+      if (!allowRefresh) {
+        throw new ControllerApplicationException(LOGGER,
+            String.format("Segment: %s already exists in table: %s. Refresh not permitted.", segmentName,
+                tableNameWithType), Response.Status.CONFLICT);
+      }
+      LOGGER.info("Segment: {} already exists in table: {}, refreshing", segmentName, tableNameWithType);
+      segmentZKMetadata.setRefreshTime(System.currentTimeMillis());
+      if (!_pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata,
+          existingZNRecord.getVersion())) {
+        throw new RuntimeException(
+            String.format("Failed to update ZK metadata for segment: %s of table: %s", segmentName,
+                tableNameWithType));
+      }
+      _pinotHelixResourceManager.sendSegmentRefreshMessage(tableNameWithType, segmentName, true, true);
+    } else {
+      if (_pinotHelixResourceManager.needTieredSegmentAssignment(tableConfig)) {
+        _pinotHelixResourceManager.updateSegmentTargetTier(tableNameWithType, segmentZKMetadata,
+            _pinotHelixResourceManager.getSortedTiers(tableConfig));
+      }
+      segmentZKMetadata.setPushTime(System.currentTimeMillis());
+      if (!_pinotHelixResourceManager.createSegmentZkMetadata(tableNameWithType, segmentZKMetadata)) {
+        throw new RuntimeException(
+            String.format("Failed to create ZK metadata for segment: %s of table: %s", segmentName,
+                tableNameWithType));
+      }
+      try {
+        _pinotHelixResourceManager.assignSegment(tableConfig, segmentZKMetadata);
+      } catch (Exception e) {
+        LOGGER.error("Failed to assign segment: {} for table: {}, ZK metadata already cleaned up by assignSegment",
+            segmentName, tableNameWithType, e);
+        throw e;
+      }
+    }


This new endpoint introduces non-trivial registration/refresh behavior (stale-ZK cleanup, create vs update, Helix assignment). The existing test suite for this class only covers helper methods; please add unit tests for registerSegmentInternal covering: new registration, refresh without wiping existing metadata, allowRefresh=false (409), and stale-ZK cleanup when IdealState is missing the segment.

Add a register endpoint to upload zk metadata directly

352aa30

Copilot AI review requested due to automatic review settings March 12, 2026 18:31

Copilot started reviewing on behalf of dinoocch March 12, 2026 18:31 View session

Copilot AI reviewed Mar 12, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add a register endpoint to upload zk metadata directly#145

Add a register endpoint to upload zk metadata directly#145
dinoocch wants to merge 1 commit into
masterfrom
dino/temp-segment-register

dinoocch commented Mar 12, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Mar 12, 2026

Uh oh!

Copilot AI Mar 12, 2026

Uh oh!

Copilot AI Mar 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		segmentZKMetadata.setRefreshTime(System.currentTimeMillis());
		if (!_pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata,

-      segmentZKMetadata.setRefreshTime(System.currentTimeMillis());
-      if (!_pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata,
+      // Merge request-provided fields into existing ZK metadata instead of overwriting the whole record.
+      SegmentZKMetadata mergedSegmentZKMetadata = new SegmentZKMetadata(existingZNRecord);
+      // Whitelist fields that are allowed to change on refresh.
+      mergedSegmentZKMetadata.setDownloadUrl(segmentZKMetadata.getDownloadUrl());
+      mergedSegmentZKMetadata.setCrc(segmentZKMetadata.getCrc());
+      mergedSegmentZKMetadata.setRefreshTime(System.currentTimeMillis());
+      if (!_pinotHelixResourceManager.updateZkMetadata(tableNameWithType, mergedSegmentZKMetadata,

-    }
+    }
+    long crc = segmentZKMetadata.getCrc();
+    if (crc < 0) {
+      throw new ControllerApplicationException(LOGGER,
+          "segment.crc is required in simpleFields and must be non-negative", Response.Status.BAD_REQUEST);
+    }

Uh oh!

Conversation

dinoocch commented Mar 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Background

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

dinoocch commented Mar 12, 2026 •

edited

Loading