Add DR and Failure scenario handling for Pauseless Ingestion

pinot-common/src/main/java/org/apache/pinot/common/metadata/segment/SegmentZKMetadataUtils.java

@@ -0,0 +1,80 @@
+/**
+ * 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.common.metadata.segment;
+
+import com.fasterxml.jackson.databind.DeserializationFeature;
+import com.fasterxml.jackson.databind.MapperFeature;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializationFeature;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import java.io.IOException;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class SegmentZKMetadataUtils {
+  private SegmentZKMetadataUtils() {
+  }
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentZKMetadataUtils.class);
+  public static final ObjectMapper MAPPER = createObjectMapper();
+
+  private static ObjectMapper createObjectMapper() {
+    ObjectMapper mapper = new ObjectMapper();
+    mapper.configure(SerializationFeature.INDENT_OUTPUT, true);
+    mapper.configure(MapperFeature.AUTO_DETECT_FIELDS, true);
+    mapper.configure(MapperFeature.AUTO_DETECT_SETTERS, true);
+    mapper.configure(MapperFeature.CAN_OVERRIDE_ACCESS_MODIFIERS, true);
+    mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
+    return mapper;
+  }
+
+  public static String serialize(SegmentZKMetadata metadata) throws IOException {
+    if (metadata == null) {
+      return null;
+    }
+    return MAPPER.writeValueAsString(metadata.toZNRecord());
+  }
+
+  public static SegmentZKMetadata deserialize(String jsonString) throws IOException {
+    if (jsonString == null || jsonString.isEmpty()) {
+      return null;
+    }
+    ObjectNode objectNode = (ObjectNode) MAPPER.readTree(jsonString);
+    ZNRecord znRecord = MAPPER.treeToValue(objectNode, ZNRecord.class);
+    return new SegmentZKMetadata(znRecord);
+  }
+
+  public static SegmentZKMetadata deserialize(ObjectNode objectNode) throws IOException {
+    if (objectNode == null) {
+      return null;
+    }
+    ZNRecord znRecord = MAPPER.treeToValue(objectNode, ZNRecord.class);
+    return new SegmentZKMetadata(znRecord);
+  }
+
+  public static SegmentZKMetadata deserialize(byte[] bytes) throws IOException {
+    if (bytes == null || bytes.length == 0) {
+      return null;
+    }
+    ZNRecord znRecord = MAPPER.readValue(bytes, ZNRecord.class);
+    return new SegmentZKMetadata(znRecord);
+  }
+}

pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java

@@ -55,9 +55,10 @@
 import org.apache.hc.core5.http.message.BasicNameValuePair;
 import org.apache.pinot.common.auth.AuthProviderUtils;
 import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.metadata.segment.SegmentZKMetadata;
+import org.apache.pinot.common.metadata.segment.SegmentZKMetadataUtils;
 import org.apache.pinot.common.restlet.resources.EndReplaceSegmentsRequest;
 import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
-import org.apache.pinot.common.restlet.resources.TableLLCSegmentUploadResponse;
 import org.apache.pinot.common.utils.http.HttpClient;
 import org.apache.pinot.common.utils.http.HttpClientConfig;
 import org.apache.pinot.spi.auth.AuthProvider;
@@ -125,6 +126,7 @@ public static FileUploadType getDefaultUploadType() {
   private static final String FORCE_CLEANUP_PARAMETER = "&forceCleanup=";
 
   private static final String RETENTION_PARAMETER = "retention=";
+  public static final String REINGEST_SEGMENT_PATH = "/reIngestSegment";
 
   private static final List<String> SUPPORTED_PROTOCOLS = Arrays.asList(HTTP, HTTPS);
 
@@ -968,26 +970,25 @@ public String uploadToSegmentStore(String uri)
    * Used by controllers to send requests to servers: Controller periodic task uses this endpoint to ask servers
    * to upload committed llc segment to segment store if missing.
    * @param uri The uri to ask servers to upload segment to segment store
-   * @return {@link TableLLCSegmentUploadResponse} - segment download url, crc, other metadata
+   * @return {@link SegmentZKMetadata} - segment download url, crc, other metadata
    * @throws URISyntaxException
    * @throws IOException
    * @throws HttpErrorStatusException
    */
-  public TableLLCSegmentUploadResponse uploadLLCToSegmentStore(String uri)
+  public SegmentZKMetadata uploadLLCToSegmentStore(String uri)
       throws URISyntaxException, IOException, HttpErrorStatusException {
     ClassicRequestBuilder requestBuilder = ClassicRequestBuilder.post(new URI(uri)).setVersion(HttpVersion.HTTP_1_1);
     // sendRequest checks the response status code
     SimpleHttpResponse response = HttpClient.wrapAndThrowHttpException(
         _httpClient.sendRequest(requestBuilder.build(), HttpClient.DEFAULT_SOCKET_TIMEOUT_MS));
-    TableLLCSegmentUploadResponse tableLLCSegmentUploadResponse = JsonUtils.stringToObject(response.getResponse(),
-        TableLLCSegmentUploadResponse.class);
-    if (tableLLCSegmentUploadResponse.getDownloadUrl() == null
-        || tableLLCSegmentUploadResponse.getDownloadUrl().isEmpty()) {
+    SegmentZKMetadata segmentZKMetadata = SegmentZKMetadataUtils.deserialize(response.getResponse());
+    if (segmentZKMetadata.getDownloadUrl() == null
+        || segmentZKMetadata.getDownloadUrl().isEmpty()) {
       throw new HttpErrorStatusException(
           String.format("Returned segment download url is empty after requesting servers to upload by the path: %s",
               uri), response.getStatusCode());
     }
-    return tableLLCSegmentUploadResponse;
+    return segmentZKMetadata;
   }
 
   /**
@@ -1248,6 +1249,63 @@ public File downloadUntarFileStreamed(URI uri, File dest, AuthProvider authProvi
         httpHeaders, maxStreamRateInByte);
   }
 
+  /**
+   * Invokes the server's reIngestSegment API via a POST request with JSON payload,
+   * using Simple HTTP APIs.
+   *
+   * POST http://[serverURL]/reIngestSegment
+   * {
+   *   "tableNameWithType": [tableName],
+   *   "segmentName": [segmentName],
+   *   "uploadURI": [leadControllerUrl],
+   *   "uploadSegment": true
+   * }
+   */
+  //TODO: Add auth and https support
+  public void triggerReIngestion(String serverHostPort, String tableNameWithType, String segmentName,
+      String leadControllerUrl)
+      throws IOException, URISyntaxException, HttpErrorStatusException {
+
+
+    if (serverHostPort.contains("http://")) {
+      serverHostPort = serverHostPort.replace("http://", "");
+    }
+
+    String serverHost = serverHostPort.split(":")[0];
+    String serverPort = serverHostPort.split(":")[1];
+
+    URI reIngestUri = getURI(HTTP, serverHost, Integer.parseInt(serverPort), REINGEST_SEGMENT_PATH);
+
+    // Build the JSON payload
+    Map<String, Object> requestJson = new HashMap<>();
+    requestJson.put("tableNameWithType", tableNameWithType);
+    requestJson.put("segmentName", segmentName);
+    requestJson.put("uploadURI", leadControllerUrl);
+    requestJson.put("uploadSegment", true);
+
+    // Convert the request payload to JSON string
+    String jsonPayload = JsonUtils.objectToString(requestJson);
+    // Prepare a POST request with Simple HTTP
+    ClassicRequestBuilder requestBuilder = ClassicRequestBuilder
+        .post(reIngestUri)
+        .setVersion(HttpVersion.HTTP_1_1)
+        .setHeader("Content-Type", "application/json")
+        .setHeader("Accept", "application/json")
+        .setEntity(new StringEntity(jsonPayload, ContentType.APPLICATION_JSON));
+
+    // Send the request using your custom HttpClient wrapper.
+    // (Adjust the timeout as needed in your environment)
+    SimpleHttpResponse response = HttpClient.wrapAndThrowHttpException(
+        _httpClient.sendRequest(requestBuilder.build(), HttpClient.DEFAULT_SOCKET_TIMEOUT_MS));
+
+    // Check that we got a 2xx response
+    int statusCode = response.getStatusCode();
+    if (statusCode / 100 != 2) {
+      throw new IOException(String.format("Failed POST to %s, HTTP %d: %s",
+          reIngestUri, statusCode, response.getResponse()));
+    }
+  }
+
   /**
    * Generate a param list with a table name attribute.
    *

pinot-common/src/test/java/org/apache/pinot/common/metadata/segment/SegmentZKMetadataUtilsTest.java

@@ -0,0 +1,126 @@
+/**
+ * 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.common.metadata.segment;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import java.io.IOException;
+import java.util.concurrent.TimeUnit;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.*;
+
+public class SegmentZKMetadataUtilsTest {
+
+  private SegmentZKMetadata _testMetadata;
+  private static final String TEST_SEGMENT_NAME = "testSegment";
+  private static final String TEST_TABLE_SEGMENT = "mytable__0__0__20220722T2342Z";
+
+  @BeforeMethod
+  public void setUp() {
+    // Create a test metadata object with sample data
+    _testMetadata = new SegmentZKMetadata(TEST_SEGMENT_NAME);
+    _testMetadata.setStartTime(1234567890L);
+    _testMetadata.setEndTime(1234567899L);
+    _testMetadata.setTimeUnit(TimeUnit.SECONDS);
+    _testMetadata.setIndexVersion("v1");
+    _testMetadata.setTotalDocs(1000);
+    _testMetadata.setSizeInBytes(1024 * 1024);
+    _testMetadata.setCrc(123456L);
+    _testMetadata.setCreationTime(System.currentTimeMillis());
+  }
+
+  @Test
+  public void testSerialize() throws IOException {
+    // Test successful serialization
+    String serialized = SegmentZKMetadataUtils.serialize(_testMetadata);
+
+    // Verify basic properties
+    assertNotNull(serialized, "Serialized string should not be null");
+    assertTrue(serialized.contains(TEST_SEGMENT_NAME), "Serialized string should contain segment name");
+    assertTrue(serialized.contains("SECONDS"), "Serialized string should contain time unit");
+
+    // Verify JSON structure
+    ObjectNode jsonNode = (ObjectNode) SegmentZKMetadataUtils.MAPPER.readTree(serialized);
+    assertTrue(jsonNode.has("simpleFields"), "Should contain simpleFields");
+    assertTrue(jsonNode.has("mapFields"), "Should contain mapFields");
+  }
+
+  @Test
+  public void testSerializeNull() throws IOException {
+    assertNull(SegmentZKMetadataUtils.serialize(null), "Serializing null should return null");
+  }
+
+  @Test
+  public void testDeserializeString() throws IOException {
+    String errorStr = "{\"id\":\"" + TEST_TABLE_SEGMENT + "\",\"simpleFields\":"
+        + "{\"segment.crc\":\"2624963047\",\"segment.creation.time\":\"1658533353347\","
+        + "\"segment.download.url\":\"http://localhost:18998/segments/mytable/" + TEST_TABLE_SEGMENT + "\","
+        + "\"segment.end.time\":\"1405296000000\",\"segment.flush.threshold.size\":\"2500\","
+        + "\"segment.index.version\":\"v3\",\"segment.realtime.endOffset\":\"2500\","
+        + "\"segment.realtime.numReplicas\":\"1\",\"segment.realtime.startOffset\":\"0\","
+        + "\"segment.realtime.status\":\"DONE\",\"segment.start.time\":\"1404086400000\","
+        + "\"segment.time.unit\":\"MILLISECONDS\",\"segment.total.docs\":\"2500\"},"
+        + "\"mapFields\":{},\"listFields\":{}}";
+
+    SegmentZKMetadata segmentZKMetadata = SegmentZKMetadataUtils.deserialize(errorStr);
+
+    // Verify deserialized properties
+    assertEquals(segmentZKMetadata.getSegmentName(), TEST_TABLE_SEGMENT,
+        "Segment name should match expected value");
+    assertEquals(segmentZKMetadata.getEndTimeMs(), 1405296000000L,
+        "End time should match expected value");
+    assertEquals(segmentZKMetadata.getCrc(), 2624963047L,
+        "CRC should match expected value");
+  }
+
+  @Test
+  public void testDeserializeObjectNode() throws IOException {
+    String errorStr = "{\"id\":\"" + TEST_TABLE_SEGMENT + "\",\"simpleFields\":"
+        + "{\"segment.crc\":\"2624963047\",\"segment.creation.time\":\"1658533353347\","
+        + "\"segment.download.url\":\"http://localhost:18998/segments/mytable/" + TEST_TABLE_SEGMENT + "\","
+        + "\"segment.end.time\":\"1405296000000\",\"segment.flush.threshold.size\":\"2500\","
+        + "\"segment.index.version\":\"v3\",\"segment.realtime.endOffset\":\"2500\","
+        + "\"segment.realtime.numReplicas\":\"1\",\"segment.realtime.startOffset\":\"0\","
+        + "\"segment.realtime.status\":\"DONE\",\"segment.start.time\":\"1404086400000\","
+        + "\"segment.time.unit\":\"MILLISECONDS\",\"segment.total.docs\":\"2500\"},"
+        + "\"mapFields\":{},\"listFields\":{}}";
+
+    JsonNode zkChildren = SegmentZKMetadataUtils.MAPPER.readTree(errorStr);
+    SegmentZKMetadata segmentZKMetadata = SegmentZKMetadataUtils.deserialize((ObjectNode) zkChildren);
+
+    assertEquals(segmentZKMetadata.getSegmentName(), TEST_TABLE_SEGMENT,
+        "Segment name should match expected value");
+    assertEquals(segmentZKMetadata.getEndTimeMs(), 1405296000000L,
+        "End time should match expected value");
+  }
+
+  @Test
+  public void testDeserializeBytes() throws IOException {
+    String serialized = SegmentZKMetadataUtils.serialize(_testMetadata);
+    byte[] bytes = serialized.getBytes();
+
+    SegmentZKMetadata deserialized = SegmentZKMetadataUtils.deserialize(bytes);
+
+    assertNotNull(deserialized, "Deserialized object should not be null");
+    assertEquals(deserialized.getSegmentName(), _testMetadata.getSegmentName(),
+        "Segment names should match");
+  }
+}

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java

@@ -405,6 +405,15 @@ public ZkHelixPropertyStore<ZNRecord> getPropertyStore() {
     return _propertyStore;
   }
 
+  /**
+   * Get the Pinot llc realtime segment manager
+   *
+   * @return Pinot llc realtime segment manager
+   */
+  public PinotLLCRealtimeSegmentManager getPinotLLCRealtimeSegmentManager() {
+    return _pinotLLCRealtimeSegmentManager;
+  }
+
   /**
    * Get the linage manager.
    *