Pauseless Consumption #3: Disaster Recovery with Reingestion

pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerGauge.java

@@ -180,7 +180,10 @@ public enum ControllerGauge implements AbstractMetrics.Gauge {
   // segment when the partition is first detected).
   COMMITTING_SEGMENT_SIZE("committingSegmentSize", false),
 
-  TABLE_REBALANCE_IN_PROGRESS("tableRebalanceInProgress", false);
+  TABLE_REBALANCE_IN_PROGRESS("tableRebalanceInProgress", false),
+
+  // Number of reingested segments getting uploaded
+  REINGESTED_SEGMENT_UPLOADS_IN_PROGRESS("reingestedSegmentUploadsInProgress", true);
 
   private final String _gaugeName;
   private final String _unit;

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

@@ -116,6 +116,7 @@ public static FileUploadType getDefaultUploadType() {
   private static final String SCHEMA_PATH = "/schemas";
   private static final String OLD_SEGMENT_PATH = "/segments";
   private static final String SEGMENT_PATH = "/v2/segments";
+  private static final String REINGESTED_SEGMENT_UPLOAD_PATH = "/segments/reingested";
   private static final String BATCH_SEGMENT_UPLOAD_PATH = "/segments/batchUpload";
   private static final String TABLES_PATH = "/tables";
   private static final String TYPE_DELIMITER = "type=";
@@ -369,6 +370,12 @@ public static URI getUploadSegmentURI(URI controllerURI)
     return getURI(controllerURI.getScheme(), controllerURI.getHost(), controllerURI.getPort(), SEGMENT_PATH);
   }
 
+  public static URI getReingestedSegmentUploadURI(URI controllerURI)
+      throws URISyntaxException {
+    return getURI(controllerURI.getScheme(), controllerURI.getHost(), controllerURI.getPort(),
+        REINGESTED_SEGMENT_UPLOAD_PATH);
+  }
+
   public static URI getBatchSegmentUploadURI(URI controllerURI)
       throws URISyntaxException {
     return getURI(controllerURI.getScheme(), controllerURI.getHost(), controllerURI.getPort(),

pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java

@@ -428,6 +428,98 @@ private SuccessResponse uploadSegment(@Nullable String tableName, TableType tabl
     }
   }
 
+  private SuccessResponse uploadReingestedSegment(String tableName, FormDataMultiPart multiPart, HttpHeaders headers,
+      Request request) {
+    tableName = DatabaseUtils.translateTableName(tableName, headers);
+    String rawTableName = TableNameBuilder.extractRawTableName(tableName);
+    String realtimeTableName = TableNameBuilder.REALTIME.tableNameWithType(tableName);
+
+    // TODO: Consider validating the segment name and table name from the header against the actual segment
+    extractHttpHeader(headers, CommonConstants.Controller.SEGMENT_NAME_HTTP_HEADER);
+    extractHttpHeader(headers, CommonConstants.Controller.TABLE_NAME_HTTP_HEADER);
+
+    String uploadTypeStr = extractHttpHeader(headers, FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE);
+    if (!FileUploadType.METADATA.name().equals(uploadTypeStr)) {
+      throw new ControllerApplicationException(LOGGER, "Reingestion upload type must be METADATA",
+          Response.Status.BAD_REQUEST);
+    }
+    String sourceDownloadURIStr = extractHttpHeader(headers, FileUploadDownloadClient.CustomHeaders.DOWNLOAD_URI);
+    if (StringUtils.isEmpty(sourceDownloadURIStr)) {
+      throw new ControllerApplicationException(LOGGER, "Source download URI is required", Response.Status.BAD_REQUEST);
+    }
+    String copySegmentToDeepStore =
+        extractHttpHeader(headers, FileUploadDownloadClient.CustomHeaders.COPY_SEGMENT_TO_DEEP_STORE);
+    if (!Boolean.parseBoolean(copySegmentToDeepStore)) {
+      throw new ControllerApplicationException(LOGGER, "COPY_SEGMENT_TO_DEEP_STORE must be true for reingestion upload",
+          Response.Status.BAD_REQUEST);
+    }
+
+    File tempTarFile = null;
+    File tempSegmentDir = null;
+    try {
+      ControllerFilePathProvider provider = ControllerFilePathProvider.getInstance();
+      String tempFileName = TMP_DIR_PREFIX + UUID.randomUUID();
+      tempTarFile = new File(provider.getFileUploadTempDir(), tempFileName);
+      tempSegmentDir = new File(provider.getUntarredFileTempDir(), tempFileName);
+
+      long segmentSizeInBytes;
+      createSegmentFileFromMultipart(multiPart, tempTarFile);
+      PinotFS pinotFS = null;
+      try {
+        URI segmentURI = new URI(sourceDownloadURIStr);
+        pinotFS = PinotFSFactory.create(segmentURI.getScheme());
+        segmentSizeInBytes = pinotFS.length(segmentURI);
+      } catch (Exception e) {
+        segmentSizeInBytes = -1;
+        LOGGER.warn("Could not fetch segment size for metadata push", e);
+      } finally {
+        if (pinotFS != null) {
+          pinotFS.close();
+        }
+      }
+
+      String metadataProviderClass = DefaultMetadataExtractor.class.getName();
+      SegmentMetadata segmentMetadata = getSegmentMetadata(tempTarFile, tempSegmentDir, metadataProviderClass);
+      String segmentName = segmentMetadata.getName();
+
+      String clientAddress = InetAddress.getByName(request.getRemoteAddr()).getHostName();
+      LOGGER.info("Processing upload request for reingested segment: {} of table: {} from client: {}", segmentName,
+          realtimeTableName, clientAddress);
+
+      // Update download URI if controller is responsible for moving the segment to the deep store
+      URI dataDirURI = provider.getDataDirURI();
+      String dataDirPath = dataDirURI.toString();
+      String encodedSegmentName = URIUtils.encode(segmentName);
+      String finalSegmentLocationPath = URIUtils.getPath(dataDirPath, rawTableName, encodedSegmentName);
+      String segmentDownloadURIStr;
+      if (dataDirURI.getScheme().equalsIgnoreCase(CommonConstants.Segment.LOCAL_SEGMENT_SCHEME)) {
+        segmentDownloadURIStr = URIUtils.getPath(provider.getVip(), "segments", rawTableName, encodedSegmentName);
+      } else {
+        segmentDownloadURIStr = finalSegmentLocationPath;
+      }
+      URI finalSegmentLocationURI = URIUtils.getUri(finalSegmentLocationPath);
+      LOGGER.info("Using segment download URI: {} for reingested segment: {} of table: {}", segmentDownloadURIStr,
+          segmentName, realtimeTableName);
+
+      ZKOperator zkOperator = new ZKOperator(_pinotHelixResourceManager, _controllerConf, _controllerMetrics);
+      zkOperator.completeReingestedSegmentOperations(realtimeTableName, segmentMetadata, finalSegmentLocationURI,
+          sourceDownloadURIStr, segmentDownloadURIStr, segmentSizeInBytes);
+
+      return new SuccessResponse(
+          "Successfully uploaded reingested segment: " + segmentName + " of table: " + realtimeTableName);
+    } 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 uploading reingested segment: " + e.getMessage(), Response.Status.INTERNAL_SERVER_ERROR, e);
+    } finally {
+      FileUtils.deleteQuietly(tempTarFile);
+      FileUtils.deleteQuietly(tempSegmentDir);
+    }
+  }
+
   // Method used to update a list of segments in batch mode with the METADATA upload type.
   private SuccessResponse uploadSegments(String tableName, TableType tableType, FormDataMultiPart multiPart,
       boolean enableParallelPushProtection, boolean allowRefresh, HttpHeaders headers, Request request) {
@@ -965,6 +1057,35 @@ public Response revertReplaceSegments(
     }
   }
 
+  @POST
+  @ManagedAsync
+  @Produces(MediaType.APPLICATION_JSON)
+  @Consumes(MediaType.MULTIPART_FORM_DATA)
+  @Path("segments/reingested")
+  @Authorize(targetType = TargetType.TABLE, paramName = "tableName", action = Actions.Table.UPLOAD_SEGMENT)
+  @Authenticate(AccessType.CREATE)
+  @ApiOperation(value = "Reingest a realtime segment", notes = "Reingest a segment as multipart file")
+  @ApiResponses(value = {
+      @ApiResponse(code = 200, message = "Successfully reingested segment"),
+      @ApiResponse(code = 400, message = "Bad Request"),
+      @ApiResponse(code = 403, message = "Segment validation fails"),
+      @ApiResponse(code = 409, message = "Segment already exists or another parallel push in progress"),
+      @ApiResponse(code = 412, message = "CRC check fails"),
+      @ApiResponse(code = 500, message = "Internal error")
+  })
+  @TrackInflightRequestMetrics
+  @TrackedByGauge(gauge = ControllerGauge.REINGESTED_SEGMENT_UPLOADS_IN_PROGRESS)
+  public void uploadReingestedSegment(FormDataMultiPart multiPart,
+      @ApiParam(value = "Name of the table", required = true)
+      @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_NAME) String tableName, @Context HttpHeaders headers,
+      @Context Request request, @Suspended AsyncResponse asyncResponse) {
+    try {
+      asyncResponse.resume(uploadReingestedSegment(tableName, multiPart, headers, request));
+    } catch (Throwable t) {
+      asyncResponse.resume(t);
+    }
+  }
+
   private static void createSegmentFileFromMultipart(FormDataMultiPart multiPart, File destFile)
       throws IOException {
     // Read segment file or segment metadata file and directly use that information to update zk

pinot-controller/src/main/java/org/apache/pinot/controller/api/upload/ZKOperator.java

@@ -43,6 +43,7 @@
 import org.apache.pinot.controller.helix.core.PinotHelixResourceManager;
 import org.apache.pinot.segment.spi.SegmentMetadata;
 import org.apache.pinot.spi.filesystem.PinotFSFactory;
+import org.apache.pinot.spi.utils.CommonConstants;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -170,6 +171,44 @@ public void completeSegmentsOperations(String tableNameWithType, FileUploadType
     processExistingSegments(tableNameWithType, uploadType, enableParallelPushProtection, headers, existingSegmentsList);
   }
 
+  public void completeReingestedSegmentOperations(String realtimeTableName, SegmentMetadata segmentMetadata,
+      URI finalSegmentLocationURI, String sourceDownloadURIStr, String segmentDownloadURIStr, long segmentSizeInBytes)
+      throws Exception {
+    String segmentName = segmentMetadata.getName();
+    ZNRecord segmentMetadataZNRecord =
+        _pinotHelixResourceManager.getSegmentMetadataZnRecord(realtimeTableName, segmentName);
+    if (segmentMetadataZNRecord == null) {
+      throw new ControllerApplicationException(LOGGER, "Failed to find segment ZK metadata for segment: " + segmentName,
+          Response.Status.NOT_FOUND);
+    }
+    int expectedVersion = segmentMetadataZNRecord.getVersion();
+    SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentMetadataZNRecord);
+    if (segmentZKMetadata.getStatus() != CommonConstants.Segment.Realtime.Status.COMMITTING) {
+      throw new ControllerApplicationException(LOGGER,
+          "Reingested segment: " + segmentName + " must be in COMMITTING status, but found: "
+              + segmentZKMetadata.getStatus(), Response.Status.CONFLICT);
+    }
+
+    // Copy the segment to the final location
+    copyFromSegmentURIToDeepStore(new URI(sourceDownloadURIStr), finalSegmentLocationURI);
+    LOGGER.info("Copied reingested segment: {} of table: {} to final location: {}", segmentName, realtimeTableName,
+        finalSegmentLocationURI);
+
+    // Update the ZK metadata
+    segmentZKMetadata.setCustomMap(segmentMetadata.getCustomMap());
+    SegmentZKMetadataUtils.updateCommittingSegmentZKMetadata(realtimeTableName, segmentZKMetadata, segmentMetadata,
+        segmentDownloadURIStr, segmentSizeInBytes, segmentZKMetadata.getEndOffset());
+    if (!_pinotHelixResourceManager.updateZkMetadata(realtimeTableName, segmentZKMetadata, expectedVersion)) {
+      throw new RuntimeException(
+          String.format("Failed to update ZK metadata for segment: %s, table: %s, expected version: %d", segmentName,
+              realtimeTableName, expectedVersion));
+    }
+    LOGGER.info("Updated reingested segment: {} of table: {} to property store", segmentName, realtimeTableName);
+
+    // Send a message to servers hosting the table to reset the segment
+    _pinotHelixResourceManager.resetSegment(realtimeTableName, segmentName, null);
+  }
+
   /**
    * Returns {@code true} when the segment should be processed as new segment.
    * <p>When segment ZK metadata exists, check if segment exists in the ideal state. If the previous upload failed after