Jackie-Jiang commented on code in PR #19180:
URL: https://github.com/apache/pinot/pull/19180#discussion_r4288923570


##########
pinot-controller/src/test/java/org/apache/pinot/controller/api/PinotSegmentRestletResourceTest.java:
##########
@@ -176,6 +183,101 @@ public void testSegmentCrcApi()
     checkCrcRequest(adminClient, rawTableName, segmentMetadataTable, 9);
   }
 
+  @Test
+  public void testUpdateSegmentZKMetadataCustomMapApi()
+      throws Exception {
+    String rawTableName = "metadataUpdateTestTable";
+    String offlineTableName = 
TableNameBuilder.OFFLINE.tableNameWithType(rawTableName);
+    String segmentName = "metadata+UpdateSegment";
+    TEST_INSTANCE.addDummySchema(rawTableName);
+    TEST_INSTANCE.getHelixResourceManager().addTable(
+        new 
TableConfigBuilder(TableType.OFFLINE).setTableName(rawTableName).setNumReplicas(1).build());
+    SegmentMetadata segmentMetadata = 
SegmentMetadataMockUtils.mockSegmentMetadata(rawTableName, segmentName);
+    TEST_INSTANCE.getHelixResourceManager().addNewSegment(offlineTableName, 
segmentMetadata, "downloadUrl");
+
+    SegmentZKMetadata originalMetadata =
+        
TEST_INSTANCE.getHelixResourceManager().getSegmentZKMetadata(offlineTableName, 
segmentName);
+    assertNotNull(originalMetadata);
+    long originalCrc = originalMetadata.getCrc();
+    long originalRefreshTime = originalMetadata.getRefreshTime();
+    URI endpoint = FileUploadDownloadClient.getUpdateSegmentZKMetadataURI(
+        URI.create(TEST_INSTANCE.getControllerBaseApiUrl()), offlineTableName, 
segmentName);
+    String modifier = new SegmentZKMetadataCustomMapModifier(
+        SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE, 
Map.of("task.lastProcessedTime", "1234"))
+        .toJsonString();
+
+    SimpleHttpResponse response = 
ControllerTest.getHttpClient().sendJsonPutRequest(endpoint, modifier,
+        Map.of(HttpHeaders.IF_MATCH, Long.toString(originalCrc)));
+    assertEquals(response.getStatusCode(), 200);
+    SegmentZKMetadata updatedMetadata =
+        
TEST_INSTANCE.getHelixResourceManager().getSegmentZKMetadata(offlineTableName, 
segmentName);
+    assertNotNull(updatedMetadata);
+    // Updating custom metadata does not replace the segment, so its refresh 
time must keep its documented meaning.
+    // Callers that need the metadata mutation time can use the ZK node mtime.
+    assertEquals(updatedMetadata.getRefreshTime(), originalRefreshTime);
+    assertEquals(updatedMetadata.getCustomMap(), 
Map.of("task.lastProcessedTime", "1234"));
+    assertEquals(updatedMetadata.getCrc(), originalCrc);
+    assertEquals(updatedMetadata.getDownloadUrl(), "downloadUrl");
+
+    response = ControllerTest.getHttpClient().sendJsonPutRequest(endpoint, 
modifier, Map.of());
+    assertEquals(response.getStatusCode(), 400);
+
+    response = ControllerTest.getHttpClient().sendJsonPutRequest(endpoint, 
"not-json",
+        Map.of(HttpHeaders.IF_MATCH, Long.toString(originalCrc)));
+    assertEquals(response.getStatusCode(), 400);
+
+    response = ControllerTest.getHttpClient().sendJsonPutRequest(endpoint, 
"x".repeat(64 * 1024 + 1),
+        Map.of(HttpHeaders.IF_MATCH, Long.toString(originalCrc)));
+    assertEquals(response.getStatusCode(), 413);
+
+    for (String malformedModifier : List.of(
+        "[]",
+        "{}",
+        "{\"mapModifyMode\":null,\"map\":{}}",
+        "{\"mapModifyMode\":1,\"map\":{}}",
+        "{\"mapModifyMode\":\"UPDATE\",\"map\":[]}",
+        "{\"mapModifyMode\":\"UPDATE\",\"map\":{\"key\":1}}",
+        
"{\"mapModifyMode\":\"UPDATE\",\"map\":{\"key\":{\"nested\":\"value\"}}}")) {
+      response = ControllerTest.getHttpClient().sendJsonPutRequest(endpoint, 
malformedModifier,
+          Map.of(HttpHeaders.IF_MATCH, Long.toString(originalCrc)));
+      assertEquals(response.getStatusCode(), 400);
+    }
+
+    response = ControllerTest.getHttpClient().sendJsonPutRequest(endpoint, 
modifier,
+        Map.of(HttpHeaders.IF_MATCH, Long.toString(originalCrc + 1)));
+    assertEquals(response.getStatusCode(), 412);
+
+    response = 
ControllerTest.getHttpClient().sendJsonPutRequest(URI.create(endpoint + 
"-missing"), modifier,

Review Comment:
   Test cleanup: endpoint already ends in /metadata, so appending "-missing" 
exercises an unknown /metadata-missing route. This returns 404 without testing 
the endpoint's missing-segment branch. Please construct the normal metadata URI 
with a nonexistent segment name instead.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java:
##########
@@ -317,6 +330,122 @@ public Map<String, Object> getSegmentMetadata(
     }
   }
 
+  @PUT
+  @Path("segments/{tableNameWithType}/{segmentName}/metadata")
+  @Authorize(targetType = TargetType.TABLE, paramName = "tableNameWithType", 
action = Actions.Table.UPLOAD_SEGMENT)
+  @Authenticate(AccessType.UPDATE)
+  @Consumes(MediaType.APPLICATION_JSON)
+  @Produces(MediaType.APPLICATION_JSON)
+  @ApiOperation(value = "Update the custom map in the ZK metadata for a 
segment",
+      notes = "Updates only the segment ZK metadata custom map without 
uploading or refreshing the segment. "
+          + "The modifier JSON request body is limited to 64 KiB.")
+  @ApiResponses(value = {
+      @ApiResponse(code = 200, message = "Success"),
+      @ApiResponse(code = 400, message = "Invalid table name, CRC, or custom 
map modifier"),
+      @ApiResponse(code = 404, message = "Table or segment not found"),
+      @ApiResponse(code = 409, message = "Segment metadata changed 
concurrently"),
+      @ApiResponse(code = 412, message = "Segment CRC does not match"),
+      @ApiResponse(code = 413, message = "Modifier exceeds 64 KiB or resulting 
ZK metadata exceeds the Helix "
+          + "serializer write limit")
+  })
+  public SuccessResponse updateSegmentZKMetadataCustomMap(
+      @ApiParam(value = "Table name with type", required = true, example = 
"myTable_OFFLINE")
+      @PathParam("tableNameWithType") String tableNameWithType,
+      @ApiParam(value = "Name of the segment", required = true) 
@PathParam("segmentName") @Encoded String segmentName,
+      @ApiParam(value = "Expected segment CRC", required = true) 
@HeaderParam(HttpHeaders.IF_MATCH)
+      String expectedCrcString,
+      @ApiParam(value = "Custom map modifier", required = true) InputStream 
customMapModifierInputStream,
+      @Context HttpHeaders headers) {
+    tableNameWithType = DatabaseUtils.translateTableName(tableNameWithType, 
headers);
+    segmentName = decodePathSegment(segmentName);
+    // The metadata path is built from these names in ZKMetadataProvider. A 
slash can address a different znode;
+    // an empty segment name cannot name a znode. Other legal segment-name 
characters are left unchanged.
+    if (tableNameWithType.contains("/") || segmentName.isEmpty() || 
segmentName.contains("/")) {
+      throw new ControllerApplicationException(LOGGER, "Invalid table or 
segment path", Status.BAD_REQUEST);
+    }
+    if (TableNameBuilder.getTableTypeFromTableName(tableNameWithType) == null) 
{
+      throw new ControllerApplicationException(LOGGER,
+          String.format("Table type not provided with table name: %s", 
tableNameWithType), Status.BAD_REQUEST);
+    }
+    String customMapModifierJson = 
readCustomMapModifier(customMapModifierInputStream);
+
+    long expectedCrc;
+    try {
+      expectedCrc = Long.parseLong(expectedCrcString);
+    } catch (Exception e) {
+      throw new ControllerApplicationException(LOGGER, "Missing or invalid 
If-Match segment CRC", Status.BAD_REQUEST,
+          e);
+    }
+
+    SegmentZKMetadataCustomMapModifier customMapModifier;
+    try {
+      customMapModifier = new 
SegmentZKMetadataCustomMapModifier(customMapModifierJson);
+    } catch (Exception e) {
+      throw new ControllerApplicationException(LOGGER, "Invalid segment ZK 
metadata custom map modifier",
+          Status.BAD_REQUEST, e);
+    }
+
+    ZNRecord segmentMetadataRecord =
+        
_pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, 
segmentName);
+    if (segmentMetadataRecord == null) {
+      throw new ControllerApplicationException(LOGGER,
+          String.format("Failed to find segment: %s in table: %s", 
segmentName, tableNameWithType), Status.NOT_FOUND);
+    }
+    SegmentZKMetadata segmentZKMetadata = new 
SegmentZKMetadata(segmentMetadataRecord);
+    if (segmentZKMetadata.getSegmentUploadStartTime() > 0) {
+      throw new ControllerApplicationException(LOGGER,
+          String.format("Segment upload is in progress for segment: %s in 
table: %s", segmentName,
+              tableNameWithType), Status.CONFLICT);
+    }
+    if (segmentZKMetadata.getCrc() != expectedCrc) {
+      throw new ControllerApplicationException(LOGGER,
+          String.format("Segment CRC does not match for segment: %s in table: 
%s", segmentName, tableNameWithType),
+          Status.PRECONDITION_FAILED);
+    }
+    
segmentZKMetadata.setCustomMap(customMapModifier.modifyMap(segmentZKMetadata.getCustomMap()));
+    boolean updated;
+    try {
+      // Let the normal Helix write serialize and enforce its 
compression-aware record size limit once.
+      updated = 
_pinotHelixResourceManager.updateZkMetadataWithoutDataChange(tableNameWithType, 
segmentZKMetadata,
+          segmentMetadataRecord.getVersion());
+    } catch (ZkMarshallingError e) {

Review Comment:
   [P2] Preserve the size failure through the actual Helix write path
   
   When a modifier under 64 KiB grows an existing record beyond the serializer 
write limit, this catch is not reached. In Helix 2.0.1, 
ZkBaseDataAccessor.doSet() catches ZkMarshallingError under catch (Exception), 
returns RetCode.ERROR, and set() returns false; ZkHelixPropertyStore inherits 
this path through ZkCacheBaseDataAccessor. Consequently this endpoint returns 
409 "metadata changed concurrently" rather than 413, encouraging retries that 
cannot succeed. The new test mocks updateZkMetadataWithoutDataChange() to throw 
directly, bypassing that behavior. Please preserve the failure reason through 
the write path or classify it after a failed write, and cover this with the 
actual Helix accessor/serializer rather than a resource-manager throw stub.



##########
pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResourceTest.java:
##########
@@ -217,6 +226,181 @@ public void 
testGetSegmentsWithInvalidPartitionMetadataWithColumnFilter()
     assertTrue(result.containsKey("seg3"));
   }
 
+  @Test
+  public void testUpdateSegmentZKMetadataCustomMapUsesNonDataChangePath() {
+    String tableNameWithType = "testTable_OFFLINE";
+    String segmentName = "testSegment";
+    SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName);
+    segmentZKMetadata.setCrc(123L);
+    ZNRecord segmentMetadataRecord = segmentZKMetadata.toZNRecord();
+    segmentMetadataRecord.setVersion(7);
+    
when(_pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, 
segmentName))
+        .thenReturn(segmentMetadataRecord);
+    
when(_pinotHelixResourceManager.updateZkMetadataWithoutDataChange(eq(tableNameWithType),
 any(), eq(7)))
+        .thenReturn(true);
+    String modifier = new SegmentZKMetadataCustomMapModifier(
+        SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE, 
Map.of("task.lastProcessedTime", "1234"))
+        .toJsonString();
+
+    
_pinotSegmentRestletResource.updateSegmentZKMetadataCustomMap(tableNameWithType,
 segmentName, "123",
+        new ByteArrayInputStream(modifier.getBytes(StandardCharsets.UTF_8)), 
null);
+
+    
verify(_pinotHelixResourceManager).updateZkMetadataWithoutDataChange(eq(tableNameWithType),
 any(), eq(7));
+    verify(_pinotHelixResourceManager, 
never()).updateZkMetadata(eq(tableNameWithType), any(), eq(7));
+  }
+
+  @Test
+  public void testUpdateSegmentZKMetadataCustomMapRejectsUploadLock() {
+    String tableNameWithType = "testTable_OFFLINE";
+    String segmentName = "testSegment";
+    SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName);
+    segmentZKMetadata.setCrc(123L);
+    segmentZKMetadata.setSegmentUploadStartTime(42L);
+    ZNRecord segmentMetadataRecord = segmentZKMetadata.toZNRecord();
+    segmentMetadataRecord.setVersion(7);
+    
when(_pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, 
segmentName))
+        .thenReturn(segmentMetadataRecord);
+    String modifier = new SegmentZKMetadataCustomMapModifier(
+        SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE, 
Map.of("task.lastProcessedTime", "1234"))
+        .toJsonString();
+
+    ControllerApplicationException exception = 
expectThrows(ControllerApplicationException.class,
+        () -> 
_pinotSegmentRestletResource.updateSegmentZKMetadataCustomMap(tableNameWithType,
 segmentName, "123",
+            new 
ByteArrayInputStream(modifier.getBytes(StandardCharsets.UTF_8)), null));
+
+    assertEquals(exception.getResponse().getStatus(), 409);
+    verify(_pinotHelixResourceManager, 
never()).updateZkMetadataWithoutDataChange(eq(tableNameWithType), any(), eq(7));
+  }
+
+  @Test
+  public void testUpdateSegmentZKMetadataCustomMapRejectsNestedPath() {
+    String modifier = new SegmentZKMetadataCustomMapModifier(
+        SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE, Map.of("key", 
"value")).toJsonString();
+    for (String segmentName : List.of("nested%2Fsegment", "")) {
+      ControllerApplicationException exception = 
expectThrows(ControllerApplicationException.class,
+          () -> 
_pinotSegmentRestletResource.updateSegmentZKMetadataCustomMap("testTable_OFFLINE",
 segmentName,
+              "123", new 
ByteArrayInputStream(modifier.getBytes(StandardCharsets.UTF_8)), null));
+      assertEquals(exception.getResponse().getStatus(), 400);
+    }
+    ControllerApplicationException invalidTable = 
expectThrows(ControllerApplicationException.class,
+        () -> 
_pinotSegmentRestletResource.updateSegmentZKMetadataCustomMap("nested/table_OFFLINE",
 "segment",
+            "123", new 
ByteArrayInputStream(modifier.getBytes(StandardCharsets.UTF_8)), null));
+    assertEquals(invalidTable.getResponse().getStatus(), 400);
+    verify(_pinotHelixResourceManager, 
never()).getSegmentMetadataZnRecord(any(), any());
+  }
+
+  @Test
+  public void testUpdateSegmentZKMetadataCustomMapDecodesPathComponent() {
+    String tableNameWithType = "testTable_OFFLINE";
+    String modifier = new SegmentZKMetadataCustomMapModifier(
+        SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE, Map.of("key", 
"value")).toJsonString();
+    Map<String, String> encodedToDecoded = Map.of(
+        "segment+plus", "segment+plus", "segment%2Bplus", "segment+plus",
+        "segment%20space", "segment space", "segment%25percent", 
"segment%percent");
+    for (Map.Entry<String, String> entry : encodedToDecoded.entrySet()) {
+      SegmentZKMetadata metadata = new SegmentZKMetadata(entry.getValue());
+      metadata.setCrc(123L);
+      ZNRecord record = metadata.toZNRecord();
+      record.setVersion(7);
+      
when(_pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, 
entry.getValue()))
+          .thenReturn(record);
+      
when(_pinotHelixResourceManager.updateZkMetadataWithoutDataChange(eq(tableNameWithType),
 any(), eq(7)))
+          .thenReturn(true);
+
+      
_pinotSegmentRestletResource.updateSegmentZKMetadataCustomMap(tableNameWithType,
 entry.getKey(), "123",
+          new ByteArrayInputStream(modifier.getBytes(StandardCharsets.UTF_8)), 
null);
+
+      
verify(_pinotHelixResourceManager).getSegmentMetadataZnRecord(tableNameWithType,
 entry.getValue());
+      clearInvocations(_pinotHelixResourceManager);

Review Comment:
   Optional cleanup: use a TestNG data provider for the encoded/decoded name 
pairs. Each invocation would get fresh mocks from @BeforeMethod, so this loop 
and clearInvocations() would be unnecessary, and failures would identify the 
individual name case.



##########
pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResourceTest.java:
##########
@@ -217,6 +226,181 @@ public void 
testGetSegmentsWithInvalidPartitionMetadataWithColumnFilter()
     assertTrue(result.containsKey("seg3"));
   }
 
+  @Test
+  public void testUpdateSegmentZKMetadataCustomMapUsesNonDataChangePath() {
+    String tableNameWithType = "testTable_OFFLINE";
+    String segmentName = "testSegment";
+    SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName);
+    segmentZKMetadata.setCrc(123L);
+    ZNRecord segmentMetadataRecord = segmentZKMetadata.toZNRecord();
+    segmentMetadataRecord.setVersion(7);
+    
when(_pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, 
segmentName))
+        .thenReturn(segmentMetadataRecord);
+    
when(_pinotHelixResourceManager.updateZkMetadataWithoutDataChange(eq(tableNameWithType),
 any(), eq(7)))
+        .thenReturn(true);
+    String modifier = new SegmentZKMetadataCustomMapModifier(
+        SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE, 
Map.of("task.lastProcessedTime", "1234"))
+        .toJsonString();
+
+    
_pinotSegmentRestletResource.updateSegmentZKMetadataCustomMap(tableNameWithType,
 segmentName, "123",
+        new ByteArrayInputStream(modifier.getBytes(StandardCharsets.UTF_8)), 
null);
+
+    
verify(_pinotHelixResourceManager).updateZkMetadataWithoutDataChange(eq(tableNameWithType),
 any(), eq(7));
+    verify(_pinotHelixResourceManager, 
never()).updateZkMetadata(eq(tableNameWithType), any(), eq(7));
+  }
+
+  @Test
+  public void testUpdateSegmentZKMetadataCustomMapRejectsUploadLock() {
+    String tableNameWithType = "testTable_OFFLINE";
+    String segmentName = "testSegment";
+    SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName);
+    segmentZKMetadata.setCrc(123L);
+    segmentZKMetadata.setSegmentUploadStartTime(42L);
+    ZNRecord segmentMetadataRecord = segmentZKMetadata.toZNRecord();
+    segmentMetadataRecord.setVersion(7);
+    
when(_pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, 
segmentName))
+        .thenReturn(segmentMetadataRecord);
+    String modifier = new SegmentZKMetadataCustomMapModifier(
+        SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE, 
Map.of("task.lastProcessedTime", "1234"))
+        .toJsonString();
+
+    ControllerApplicationException exception = 
expectThrows(ControllerApplicationException.class,
+        () -> 
_pinotSegmentRestletResource.updateSegmentZKMetadataCustomMap(tableNameWithType,
 segmentName, "123",
+            new 
ByteArrayInputStream(modifier.getBytes(StandardCharsets.UTF_8)), null));
+
+    assertEquals(exception.getResponse().getStatus(), 409);
+    verify(_pinotHelixResourceManager, 
never()).updateZkMetadataWithoutDataChange(eq(tableNameWithType), any(), eq(7));
+  }
+
+  @Test
+  public void testUpdateSegmentZKMetadataCustomMapRejectsNestedPath() {
+    String modifier = new SegmentZKMetadataCustomMapModifier(
+        SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE, Map.of("key", 
"value")).toJsonString();
+    for (String segmentName : List.of("nested%2Fsegment", "")) {
+      ControllerApplicationException exception = 
expectThrows(ControllerApplicationException.class,
+          () -> 
_pinotSegmentRestletResource.updateSegmentZKMetadataCustomMap("testTable_OFFLINE",
 segmentName,
+              "123", new 
ByteArrayInputStream(modifier.getBytes(StandardCharsets.UTF_8)), null));
+      assertEquals(exception.getResponse().getStatus(), 400);
+    }
+    ControllerApplicationException invalidTable = 
expectThrows(ControllerApplicationException.class,
+        () -> 
_pinotSegmentRestletResource.updateSegmentZKMetadataCustomMap("nested/table_OFFLINE",
 "segment",
+            "123", new 
ByteArrayInputStream(modifier.getBytes(StandardCharsets.UTF_8)), null));
+    assertEquals(invalidTable.getResponse().getStatus(), 400);
+    verify(_pinotHelixResourceManager, 
never()).getSegmentMetadataZnRecord(any(), any());
+  }
+
+  @Test
+  public void testUpdateSegmentZKMetadataCustomMapDecodesPathComponent() {
+    String tableNameWithType = "testTable_OFFLINE";
+    String modifier = new SegmentZKMetadataCustomMapModifier(
+        SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE, Map.of("key", 
"value")).toJsonString();
+    Map<String, String> encodedToDecoded = Map.of(
+        "segment+plus", "segment+plus", "segment%2Bplus", "segment+plus",
+        "segment%20space", "segment space", "segment%25percent", 
"segment%percent");
+    for (Map.Entry<String, String> entry : encodedToDecoded.entrySet()) {
+      SegmentZKMetadata metadata = new SegmentZKMetadata(entry.getValue());
+      metadata.setCrc(123L);
+      ZNRecord record = metadata.toZNRecord();
+      record.setVersion(7);
+      
when(_pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, 
entry.getValue()))
+          .thenReturn(record);
+      
when(_pinotHelixResourceManager.updateZkMetadataWithoutDataChange(eq(tableNameWithType),
 any(), eq(7)))
+          .thenReturn(true);
+
+      
_pinotSegmentRestletResource.updateSegmentZKMetadataCustomMap(tableNameWithType,
 entry.getKey(), "123",
+          new ByteArrayInputStream(modifier.getBytes(StandardCharsets.UTF_8)), 
null);
+
+      
verify(_pinotHelixResourceManager).getSegmentMetadataZnRecord(tableNameWithType,
 entry.getValue());
+      clearInvocations(_pinotHelixResourceManager);
+    }
+  }
+
+  @Test
+  public void testUpdateSegmentZKMetadataCustomMapRejectsOversizedModifier() {
+    String modifier = new SegmentZKMetadataCustomMapModifier(
+        SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE, Map.of("key", 
"x".repeat(64 * 1024)))
+        .toJsonString();
+
+    ControllerApplicationException exception = 
expectThrows(ControllerApplicationException.class,
+        () -> 
_pinotSegmentRestletResource.updateSegmentZKMetadataCustomMap("testTable_OFFLINE",
 "segment", "123",
+            new 
ByteArrayInputStream(modifier.getBytes(StandardCharsets.UTF_8)), null));
+
+    assertEquals(exception.getResponse().getStatus(), 413);
+    verify(_pinotHelixResourceManager, 
never()).getSegmentMetadataZnRecord(any(), any());
+  }
+
+  @Test
+  public void 
testUpdateSegmentZKMetadataCustomMapRejectsConcurrentLockAcquisition() {
+    String tableNameWithType = "testTable_OFFLINE";
+    String segmentName = "testSegment";
+    SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName);
+    segmentZKMetadata.setCrc(123L);
+    ZNRecord segmentMetadataRecord = segmentZKMetadata.toZNRecord();
+    segmentMetadataRecord.setVersion(7);
+    
when(_pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, 
segmentName))
+        .thenReturn(segmentMetadataRecord);
+    // The uploader acquires its lock after this request reads the record, 
causing the versioned write to fail.
+    
when(_pinotHelixResourceManager.updateZkMetadataWithoutDataChange(eq(tableNameWithType),
 any(), eq(7)))
+        .thenReturn(false);
+    String modifier = new SegmentZKMetadataCustomMapModifier(
+        SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE, 
Map.of("task.lastProcessedTime", "1234"))
+        .toJsonString();
+
+    ControllerApplicationException exception = 
expectThrows(ControllerApplicationException.class,
+        () -> 
_pinotSegmentRestletResource.updateSegmentZKMetadataCustomMap(tableNameWithType,
 segmentName, "123",
+            new 
ByteArrayInputStream(modifier.getBytes(StandardCharsets.UTF_8)), null));
+
+    assertEquals(exception.getResponse().getStatus(), 409);
+    
verify(_pinotHelixResourceManager).updateZkMetadataWithoutDataChange(eq(tableNameWithType),
 any(), eq(7));
+  }
+
+  @Test
+  public void testUpdateSegmentZKMetadataCustomMapDoesNotAddRecordSizeLimit() {
+    String tableNameWithType = "testTable_OFFLINE";
+    String segmentName = "testSegment";
+    SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName);
+    segmentZKMetadata.setCrc(123L);
+    segmentZKMetadata.getSimpleFields().put("oversized", "x".repeat(950_000));

Review Comment:
   Optional test cleanup: this creates a 950 KB string but stubs the write to 
succeed and only verifies that the helper was called. It adds little beyond the 
happy-path test and does not establish real serializer-limit behavior. Consider 
folding it into the actual Helix size-boundary coverage requested above, with 
an accepted record and an over-limit record, instead of retaining this 
artificial large-record case.



##########
pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java:
##########
@@ -317,6 +330,122 @@ public Map<String, Object> getSegmentMetadata(
     }
   }
 
+  @PUT
+  @Path("segments/{tableNameWithType}/{segmentName}/metadata")
+  @Authorize(targetType = TargetType.TABLE, paramName = "tableNameWithType", 
action = Actions.Table.UPLOAD_SEGMENT)
+  @Authenticate(AccessType.UPDATE)
+  @Consumes(MediaType.APPLICATION_JSON)
+  @Produces(MediaType.APPLICATION_JSON)
+  @ApiOperation(value = "Update the custom map in the ZK metadata for a 
segment",
+      notes = "Updates only the segment ZK metadata custom map without 
uploading or refreshing the segment. "
+          + "The modifier JSON request body is limited to 64 KiB.")
+  @ApiResponses(value = {
+      @ApiResponse(code = 200, message = "Success"),
+      @ApiResponse(code = 400, message = "Invalid table name, CRC, or custom 
map modifier"),
+      @ApiResponse(code = 404, message = "Table or segment not found"),
+      @ApiResponse(code = 409, message = "Segment metadata changed 
concurrently"),
+      @ApiResponse(code = 412, message = "Segment CRC does not match"),
+      @ApiResponse(code = 413, message = "Modifier exceeds 64 KiB or resulting 
ZK metadata exceeds the Helix "
+          + "serializer write limit")
+  })
+  public SuccessResponse updateSegmentZKMetadataCustomMap(
+      @ApiParam(value = "Table name with type", required = true, example = 
"myTable_OFFLINE")
+      @PathParam("tableNameWithType") String tableNameWithType,
+      @ApiParam(value = "Name of the segment", required = true) 
@PathParam("segmentName") @Encoded String segmentName,
+      @ApiParam(value = "Expected segment CRC", required = true) 
@HeaderParam(HttpHeaders.IF_MATCH)
+      String expectedCrcString,
+      @ApiParam(value = "Custom map modifier", required = true) InputStream 
customMapModifierInputStream,
+      @Context HttpHeaders headers) {
+    tableNameWithType = DatabaseUtils.translateTableName(tableNameWithType, 
headers);
+    segmentName = decodePathSegment(segmentName);
+    // The metadata path is built from these names in ZKMetadataProvider. A 
slash can address a different znode;
+    // an empty segment name cannot name a znode. Other legal segment-name 
characters are left unchanged.
+    if (tableNameWithType.contains("/") || segmentName.isEmpty() || 
segmentName.contains("/")) {
+      throw new ControllerApplicationException(LOGGER, "Invalid table or 
segment path", Status.BAD_REQUEST);
+    }
+    if (TableNameBuilder.getTableTypeFromTableName(tableNameWithType) == null) 
{
+      throw new ControllerApplicationException(LOGGER,
+          String.format("Table type not provided with table name: %s", 
tableNameWithType), Status.BAD_REQUEST);
+    }
+    String customMapModifierJson = 
readCustomMapModifier(customMapModifierInputStream);
+
+    long expectedCrc;
+    try {
+      expectedCrc = Long.parseLong(expectedCrcString);
+    } catch (Exception e) {

Review Comment:
   Optional cleanup: narrow this catch to NumberFormatException. 
Long.parseLong(String) already throws that exception for both a missing/null 
value and malformed input, so catch (Exception) is broader than needed here.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to