You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by GitBox <gi...@apache.org> on 2022/08/30 00:47:23 UTC

[GitHub] [pinot] yuanbenson opened a new pull request, #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

yuanbenson opened a new pull request, #9295:
URL: https://github.com/apache/pinot/pull/9295

   ### Description: 
   This PR addresses https://github.com/apache/pinot/issues/9268 for the segment push job runners under the standalone execution framework: `SegmentMetadataPushJobRunner`, `SegmentTarPushJobRunner`, and `SegmentUriPushJobRunner`. 
   
   This is accomplished by introducing a new class `ConsistentDataPushUtils` which contains APIs and helpers for `*PushJobRunner(s)` to call to invoke the consistent push protocol. 
   
   Since there are large overlaps in the code for all of the `*PushJobRunner(s)`, also took this opportunity to refactor and extract the common logic out to `BaseSegmentPushJobRunner` abstract class.
   
   To enable consistent data push, this PR also introduces a new boolean config in table config under 
   `TableConfig->IngestionConfig->BatchIngestionConfig->consistentDataPush`.
   
   Users can enable consistent data push by setting the `consistentDataPush` config to true as below before invoking ingestion jobs,
   ```
   ...
       "batchIngestionConfig": {
         "segmentIngestionType": "REFRESH",
         "segmentIngestionFrequency": "DAILY",
         "consistentDataPush": true
       },
   ...
   ```
   which will 
   1) _**In the segment generation phase**_: inject timestamps to segment names (via segment postfix) in order to prevent segment name conflicts from overwriting existing segments directly for the REFRESH usecase and 
   2) _**In the segment push/upload phase**_: wrap the segments replacement protocol around segments upload which will achieve atomic switching between old and new segments data. In the case of failure, revert and abort the swap, which will ensure broker only routes to the old segments. 
   
   
   ### Testing Done: 
   Added new test `testUploadAndQueryWithConsistentPush` in `SegmentUploadIntegrationTest`, which 
   1. Runs SegmentMetadataPushJobRunner with consistent push enabled.
   [] -> [v1 segments]
   2. Checks that the segment lineage entry is in expected and completed state.
   3. Checks that count stars return expected outputs. 
   
   4. Runs SegmentTarPushJobRunner with consistent push enabled.
   [v1 segments] -> [v2 segments]
   5. Checks again that the segment lineage entry is in expected and completed state.
   6. Checks again that count stars return expected outputs (that we have successfully bulk replaced the original set of segments). 
   
   


-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] Jackie-Jiang commented on pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
Jackie-Jiang commented on PR #9295:
URL: https://github.com/apache/pinot/pull/9295#issuecomment-1239909487

   This is a great feature! Can you please help also update the pinot documentation about this new feature? https://github.com/pinot-contrib/pinot-docs


-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r958994759


##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentUploadIntegrationTest.java:
##########
@@ -81,10 +89,15 @@ protected List<String> getBloomFilterColumns() {
     return null;
   }
 
+  @BeforeMethod
+  public void setUpTest()

Review Comment:
   Calling `TestUtils.ensureDirectoriesExistAndEmpty(_tempDir, _segmentDir, _tarDir)` in `setUp` is insufficient as there is another test case `testUploadAndQuery` that leaves data in the tmp directories.
   
   Using two table names will also lead to more code for now, since it seems like `BaseClusterIntegrationTest` is written in a way that assumes one table/config per subclass, e.g. `getTableName`, `createOfflineTableConfig`. 
   
   Hence, I prefer to keep the test structure as is with the before and after method annotations, unless this is an anti-pattern. Another way is to create a different test class altogether, but combining the two here to not prevent slow down from too many tests. 



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson closed pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson closed pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners
URL: https://github.com/apache/pinot/pull/9295


-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r964301363


##########
pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-standalone/src/main/java/org/apache/pinot/plugin/ingestion/batch/standalone/SegmentTarPushJobRunner.java:
##########
@@ -46,48 +35,8 @@ public SegmentTarPushJobRunner(SegmentGenerationJobSpec spec) {
     init(spec);
   }
 
-  @Override
-  public void init(SegmentGenerationJobSpec spec) {
-    _spec = spec;
-  }
-
-  @Override
-  public void run() {
-    //init all file systems
-    List<PinotFSSpec> pinotFSSpecs = _spec.getPinotFSSpecs();
-    for (PinotFSSpec pinotFSSpec : pinotFSSpecs) {
-      PinotFSFactory.register(pinotFSSpec.getScheme(), pinotFSSpec.getClassName(), new PinotConfiguration(pinotFSSpec));
-    }
-
-    //Get outputFS for writing output pinot segments
-    URI outputDirURI;
-    try {
-      outputDirURI = new URI(_spec.getOutputDirURI());
-      if (outputDirURI.getScheme() == null) {
-        outputDirURI = new File(_spec.getOutputDirURI()).toURI();
-      }
-    } catch (URISyntaxException e) {
-      throw new RuntimeException("outputDirURI is not valid - '" + _spec.getOutputDirURI() + "'");
-    }
-    PinotFS outputDirFS = PinotFSFactory.create(outputDirURI.getScheme());
-    //Get list of files to process
-    String[] files;
-    try {
-      files = outputDirFS.listFiles(outputDirURI, true);
-    } catch (IOException e) {
-      throw new RuntimeException("Unable to list all files under outputDirURI - '" + outputDirURI + "'");
-    }
-
-    List<String> segmentsToPush = new ArrayList<>();
-    for (String file : files) {
-      if (file.endsWith(Constants.TAR_GZ_FILE_EXT)) {
-        segmentsToPush.add(file);
-      }
-    }
-    try {
-      SegmentPushUtils.pushSegments(_spec, outputDirFS, segmentsToPush);
-    } catch (RetriableOperationException | AttemptsExceededException e) {
-      throw new RuntimeException(e);
-    }
+  public void uploadSegments(Map<String, String> segmentsUriToTarPathMap)
+      throws AttemptsExceededException, RetriableOperationException {
+    SegmentPushUtils.pushSegments(_spec, _outputDirFS, new ArrayList<>(segmentsUriToTarPathMap.values()));

Review Comment:
   @snleee Do you think this might cause some performance loss (in the case of having many segments to push)? If so, might need to change `pushSegments` to take in `Collection<String>` as opposed to `List<String>`.



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r959045654


##########
pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java:
##########
@@ -747,6 +751,54 @@ public SimpleHttpResponse uploadSegment(URI uri, String segmentName, InputStream
     return uploadSegment(uri, segmentName, inputStream, null, parameters, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
   }
 
+  /**
+   * Returns a map from a given tableType to a list of segments for that given tableType (OFFLINE or REALTIME)
+   * If tableType is left unspecified, both OFFLINE and REALTIME segments will be returned in the map.
+   */
+  public Map<String, List<String>> getSegments(URI uri, String rawTableName, @Nullable String tableType,
+      boolean excludeReplacedSegments)
+      throws URISyntaxException, IOException {
+    List<String> tableTypes;
+    if (tableType == null || tableType.isEmpty()) {
+      tableTypes = Arrays.asList(TableType.OFFLINE.toString(), TableType.REALTIME.toString());
+    } else {
+      tableTypes = Arrays.asList(tableType);
+    }
+    ControllerRequestURLBuilder controllerRequestURLBuilder = ControllerRequestURLBuilder.baseUrl(uri.toString());
+    Map<String, List<String>> tableTypeToSegments = new HashMap<>();
+    for (String tableTypeToFilter : tableTypes) {
+      List<String> segments = new ArrayList<>();
+      RequestBuilder requestBuilder = RequestBuilder.get(
+          controllerRequestURLBuilder.forSegmentListAPIWithTableTypeAndExcludeReplacedSegments(rawTableName,
+              tableTypeToFilter, excludeReplacedSegments)).setVersion(HttpVersion.HTTP_1_1);
+      HttpClient.setTimeout(requestBuilder, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
+      SimpleHttpResponse response;
+      try {
+        response = HttpClient.wrapAndThrowHttpException(_httpClient.sendRequest(requestBuilder.build()));
+      } catch (HttpErrorStatusException e) {
+        tableTypeToSegments.put(tableTypeToFilter, segments);

Review Comment:
   Addressed! Logging at INFO level (as opposed to WARN) here since this catch clause will be triggered when users pass in null as `tableType`, which is an intended feature. 



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r963981920


##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java:
##########
@@ -0,0 +1,285 @@
+/**
+ * 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.segment.local.utils;
+
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.common.auth.AuthProviderUtils;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.apache.pinot.spi.utils.retry.RetryPolicy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class ConsistentDataPushUtils {
+  private ConsistentDataPushUtils() {
+  }
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentPushUtils.class);
+  private static final FileUploadDownloadClient FILE_UPLOAD_DOWNLOAD_CLIENT = new FileUploadDownloadClient();
+  private static final RetryPolicy DEFAULT_RETRY_POLICY = RetryPolicies.exponentialBackoffRetryPolicy(5, 10_000L, 2.0);
+  public static final String SEGMENT_NAME_POSTFIX = "segment.name.postfix";
+
+  /**
+   * Checks for enablement of consistent data push. If enabled, fetch the list of segments to be replaced, then
+   * invoke startReplaceSegments API and returns a map of controller URI to lineage entry IDs.
+   * If not, returns an empty hashmap.
+   */
+  public static Map<URI, String> preUpload(SegmentGenerationJobSpec spec, List<String> segmentsTo)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    boolean consistentDataPushEnabled = consistentDataPushEnabled(spec);

Review Comment:
   We already check if the consistent push is enabled in the caller side. We can remove the check for consistent push.



##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java:
##########
@@ -0,0 +1,285 @@
+/**
+ * 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.segment.local.utils;
+
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.common.auth.AuthProviderUtils;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.apache.pinot.spi.utils.retry.RetryPolicy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class ConsistentDataPushUtils {
+  private ConsistentDataPushUtils() {
+  }
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentPushUtils.class);
+  private static final FileUploadDownloadClient FILE_UPLOAD_DOWNLOAD_CLIENT = new FileUploadDownloadClient();
+  private static final RetryPolicy DEFAULT_RETRY_POLICY = RetryPolicies.exponentialBackoffRetryPolicy(5, 10_000L, 2.0);
+  public static final String SEGMENT_NAME_POSTFIX = "segment.name.postfix";
+
+  /**
+   * Checks for enablement of consistent data push. If enabled, fetch the list of segments to be replaced, then
+   * invoke startReplaceSegments API and returns a map of controller URI to lineage entry IDs.
+   * If not, returns an empty hashmap.
+   */
+  public static Map<URI, String> preUpload(SegmentGenerationJobSpec spec, List<String> segmentsTo)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    boolean consistentDataPushEnabled = consistentDataPushEnabled(spec);
+    LOGGER.info("Consistent data push is: {}", consistentDataPushEnabled ? "enabled" : "disabled");
+    Map<URI, String> uriToLineageEntryIdMap = null;
+    if (consistentDataPushEnabled) {
+      LOGGER.info("Start consistent push for table: " + rawTableName);
+      Map<URI, List<String>> uriToExistingOfflineSegments = getSegmentsToReplace(spec, rawTableName);
+      LOGGER.info("Existing segments for table {}: " + uriToExistingOfflineSegments, rawTableName);
+      LOGGER.info("New segments for table: {}: " + segmentsTo, rawTableName);
+      uriToLineageEntryIdMap = startReplaceSegments(spec, uriToExistingOfflineSegments, segmentsTo);
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * uriToLineageEntryIdMap is non-empty if and only if consistent data push is enabled.
+   * If uriToLineageEntryIdMap is non-empty, end the consistent data push protocol for each controller.
+   */
+  public static void postUpload(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    if (uriToLineageEntryIdMap != null && !uriToLineageEntryIdMap.isEmpty()) {
+      LOGGER.info("End consistent push for table: " + rawTableName);
+      endReplaceSegments(spec, uriToLineageEntryIdMap);
+    }
+  }
+
+  /**
+   * Builds a map of controller URI to startReplaceSegments URI for each Pinot cluster in the spec.
+   */
+  public static Map<URI, URI> getStartReplaceSegmentUris(SegmentGenerationJobSpec spec, String rawTableName) {
+    Map<URI, URI> baseUriToStartReplaceSegmentUriMap = new HashMap<>();
+    for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
+      URI controllerURI;
+      try {
+        controllerURI = new URI(pinotClusterSpec.getControllerURI());
+        baseUriToStartReplaceSegmentUriMap.put(controllerURI,
+            FileUploadDownloadClient.getStartReplaceSegmentsURI(controllerURI, rawTableName,
+                TableType.OFFLINE.toString(), true));
+      } catch (URISyntaxException e) {
+        throw new RuntimeException("Got invalid controller uri - '" + pinotClusterSpec.getControllerURI() + "'");
+      }
+    }
+    return baseUriToStartReplaceSegmentUriMap;
+  }
+
+  /**
+   * Starts consistent data push protocol for each Pinot cluster in the spec.
+   * Returns a map of controller URI to segment lineage entry ID.
+   */
+  public static Map<URI, String> startReplaceSegments(SegmentGenerationJobSpec spec,
+      Map<URI, List<String>> uriToSegmentsFrom, List<String> segmentsTo)
+      throws Exception {
+    Map<URI, String> uriToLineageEntryIdMap = new HashMap<>();
+    String rawTableName = spec.getTableSpec().getTableName();
+    Map<URI, URI> segmentsUris = getStartReplaceSegmentUris(spec, rawTableName);
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    LOGGER.info("Start replace segment URIs: " + segmentsUris);
+
+    for (Map.Entry<URI, URI> entry : segmentsUris.entrySet()) {
+      URI controllerUri = entry.getKey();
+      URI startSegmentUri = entry.getValue();
+      List<String> segmentsFrom = uriToSegmentsFrom.get(controllerUri);
+
+      StartReplaceSegmentsRequest startReplaceSegmentsRequest =
+          new StartReplaceSegmentsRequest(segmentsFrom, segmentsTo);
+      DEFAULT_RETRY_POLICY.attempt(() -> {
+        try {
+          SimpleHttpResponse response =
+              FILE_UPLOAD_DOWNLOAD_CLIENT.startReplaceSegments(startSegmentUri, startReplaceSegmentsRequest,
+                  authProvider);
+          String responseString = response.getResponse();
+          LOGGER.info(
+              "Got response {}: {} while sending start replace segment request for table: {}, uploadURI: {}, request:"
+                  + " {}", response.getStatusCode(), responseString, rawTableName, startSegmentUri,
+              startReplaceSegmentsRequest);
+          String segmentLineageEntryId =
+              JsonUtils.stringToJsonNode(responseString).get("segmentLineageEntryId").asText();
+          uriToLineageEntryIdMap.put(controllerUri, segmentLineageEntryId);
+          return true;
+        } catch (SocketTimeoutException se) {
+          // In case of the timeout, we should re-try.
+          return false;
+        } catch (HttpErrorStatusException e) {
+          if (e.getStatusCode() >= 500) {
+            return false;
+          } else {
+            if (e.getStatusCode() == Response.Status.NOT_FOUND.getStatusCode()) {
+              LOGGER.error("Table: {} not found when sending request: {}", rawTableName, startSegmentUri);
+            }
+            throw e;
+          }
+        }
+      });
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * Ends consistent data push protocol for each Pinot cluster in the spec.
+   */
+  public static void endReplaceSegments(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap)
+      throws Exception {
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    String rawTableName = spec.getTableSpec().getTableName();
+    for (URI controllerUri : uriToLineageEntryIdMap.keySet()) {
+      String segmentLineageEntryId = uriToLineageEntryIdMap.get(controllerUri);
+      URI uri =
+          FileUploadDownloadClient.getEndReplaceSegmentsURI(controllerUri, rawTableName, TableType.OFFLINE.toString(),
+              segmentLineageEntryId);
+      DEFAULT_RETRY_POLICY.attempt(() -> {
+        try {
+          SimpleHttpResponse response =
+              FILE_UPLOAD_DOWNLOAD_CLIENT.endReplaceSegments(uri, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS, authProvider);
+          LOGGER.info("Got response {}: {} while sending end replace segment request for table: {}, uploadURI: {}",
+              response.getStatusCode(), response.getResponse(), rawTableName, uri);
+          return true;
+        } catch (SocketTimeoutException se) {
+          // In case of the timeout, we should re-try.
+          return false;
+        } catch (HttpErrorStatusException e) {
+          if (e.getStatusCode() >= 500) {
+            return false;
+          } else {
+            throw e;
+          }
+        }
+      });
+    }
+  }
+
+  /**
+   * Revert segment lineage entry when exception gets caught. This revert request is called at best effort.
+   * If the revert call fails at this point, the next startReplaceSegment call will do the cleanup
+   * by marking the previous entry to "REVERTED" and cleaning up the leftover segments.
+   */
+  public static void handleUploadException(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap,
+      Exception exception) {
+    if (uriToLineageEntryIdMap != null) {
+      LOGGER.error("Exception when pushing segments. Marking segment lineage entry to 'REVERTED'.", exception);
+      String rawTableName = spec.getTableSpec().getTableName();
+      for (Map.Entry<URI, String> entry : uriToLineageEntryIdMap.entrySet()) {
+        String segmentLineageEntryId = entry.getValue();
+        try {
+          URI uri = FileUploadDownloadClient.getRevertReplaceSegmentsURI(entry.getKey(), rawTableName,
+              TableType.OFFLINE.name(), segmentLineageEntryId, true);
+          SimpleHttpResponse response = FILE_UPLOAD_DOWNLOAD_CLIENT.revertReplaceSegments(uri);
+          LOGGER.info("Got response {}: {} while sending revert replace segment request for table: {}, uploadURI: {}",
+              response.getStatusCode(), response.getResponse(), rawTableName, entry.getKey());
+        } catch (URISyntaxException | HttpErrorStatusException | IOException e) {
+          LOGGER.error("Exception when sending revert replace segment request to controller: {} for table: {}",
+              entry.getKey(), rawTableName, e);
+        }
+      }
+    }
+  }
+
+  public static TableConfig getTableConfig(SegmentGenerationJobSpec spec) {
+    String rawTableName = spec.getTableSpec().getTableName();
+    TableConfig tableConfig =
+        SegmentGenerationUtils.getTableConfig(spec.getTableSpec().getTableConfigURI(), spec.getAuthToken());
+    Preconditions.checkState(tableConfig != null, "Failed to find table config for table: %s", rawTableName);
+    return tableConfig;
+  }
+
+  public static boolean consistentDataPushEnabled(SegmentGenerationJobSpec spec) {
+    TableConfig tableConfig = getTableConfig(spec);
+    // Enable consistent data push only if "consistentDataPush" is set to true in batch ingestion config and the
+    // table is REFRESH use case.
+    // TODO: Remove the check for REFRESH when we support consistent push for APPEND table
+    return "REFRESH".equalsIgnoreCase(IngestionConfigUtils.getBatchSegmentIngestionType(tableConfig))
+        && IngestionConfigUtils.getBatchSegmentIngestionConsistentDataPushEnabled(tableConfig);
+  }
+
+  /**
+   * Returns a map of controller URI to a list of existing OFFLINE segments.
+   */
+  public static Map<URI, List<String>> getSegmentsToReplace(SegmentGenerationJobSpec spec, String rawTableName) {
+    Map<URI, List<String>> uriToOfflineSegments = new HashMap<>();
+    for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
+      URI controllerURI;
+      List<String> offlineSegments;
+      try {
+        controllerURI = new URI(pinotClusterSpec.getControllerURI());
+        Map<String, List<String>> segments =
+            FILE_UPLOAD_DOWNLOAD_CLIENT.getSegments(controllerURI, rawTableName, TableType.OFFLINE, true);
+        offlineSegments = segments.get(TableType.OFFLINE.toString());
+        uriToOfflineSegments.put(controllerURI, offlineSegments);
+      } catch (URISyntaxException e) {
+        throw new RuntimeException("Got invalid controller uri - '" + pinotClusterSpec.getControllerURI() + "'");
+      } catch (Exception e) {
+        e.printStackTrace();
+      }
+    }
+    return uriToOfflineSegments;
+  }
+
+  /**
+   * If consistent data push is enabled, append current timestamp to existing configured segment name postfix, if
+   * configured, to make segment name unique.
+   */
+  public static void configureSegmentPostfix(SegmentGenerationJobSpec spec) {
+    SegmentNameGeneratorSpec segmentNameGeneratorSpec = spec.getSegmentNameGeneratorSpec();
+    if (consistentDataPushEnabled(spec)) {

Review Comment:
   From the `SegmentGenerationJobRunner:init()` call, we already fetch table config. We should set `_consistentDataPushEnabled` there instead of fetching the table config again here.
   
   For all cases, we should read table config only once from the controller. (once for segment generation, once for segment push)



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r964338431


##########
pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-standalone/src/main/java/org/apache/pinot/plugin/ingestion/batch/standalone/SegmentTarPushJobRunner.java:
##########
@@ -46,48 +35,8 @@ public SegmentTarPushJobRunner(SegmentGenerationJobSpec spec) {
     init(spec);
   }
 
-  @Override
-  public void init(SegmentGenerationJobSpec spec) {
-    _spec = spec;
-  }
-
-  @Override
-  public void run() {
-    //init all file systems
-    List<PinotFSSpec> pinotFSSpecs = _spec.getPinotFSSpecs();
-    for (PinotFSSpec pinotFSSpec : pinotFSSpecs) {
-      PinotFSFactory.register(pinotFSSpec.getScheme(), pinotFSSpec.getClassName(), new PinotConfiguration(pinotFSSpec));
-    }
-
-    //Get outputFS for writing output pinot segments
-    URI outputDirURI;
-    try {
-      outputDirURI = new URI(_spec.getOutputDirURI());
-      if (outputDirURI.getScheme() == null) {
-        outputDirURI = new File(_spec.getOutputDirURI()).toURI();
-      }
-    } catch (URISyntaxException e) {
-      throw new RuntimeException("outputDirURI is not valid - '" + _spec.getOutputDirURI() + "'");
-    }
-    PinotFS outputDirFS = PinotFSFactory.create(outputDirURI.getScheme());
-    //Get list of files to process
-    String[] files;
-    try {
-      files = outputDirFS.listFiles(outputDirURI, true);
-    } catch (IOException e) {
-      throw new RuntimeException("Unable to list all files under outputDirURI - '" + outputDirURI + "'");
-    }
-
-    List<String> segmentsToPush = new ArrayList<>();
-    for (String file : files) {
-      if (file.endsWith(Constants.TAR_GZ_FILE_EXT)) {
-        segmentsToPush.add(file);
-      }
-    }
-    try {
-      SegmentPushUtils.pushSegments(_spec, outputDirFS, segmentsToPush);
-    } catch (RetriableOperationException | AttemptsExceededException e) {
-      throw new RuntimeException(e);
-    }
+  public void uploadSegments(Map<String, String> segmentsUriToTarPathMap)
+      throws AttemptsExceededException, RetriableOperationException {
+    SegmentPushUtils.pushSegments(_spec, _outputDirFS, new ArrayList<>(segmentsUriToTarPathMap.values()));

Review Comment:
   For copying the elements from `map.values()` to `List`? 
   
   In our case, the worst-case scenario should be 10s thousands of segments for a table and I think it's OK to keep it this way. I don't think that it's a performance-critical path.



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] codecov-commenter commented on pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
codecov-commenter commented on PR #9295:
URL: https://github.com/apache/pinot/pull/9295#issuecomment-1232272367

   # [Codecov](https://codecov.io/gh/apache/pinot/pull/9295?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) Report
   > Merging [#9295](https://codecov.io/gh/apache/pinot/pull/9295?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) (27443ab) into [master](https://codecov.io/gh/apache/pinot/commit/5ecca80cc8ac763d5c3ee256ea9504f60c581ecd?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) (5ecca80) will **increase** coverage by `13.10%`.
   > The diff coverage is `11.24%`.
   
   ```diff
   @@              Coverage Diff              @@
   ##             master    #9295       +/-   ##
   =============================================
   + Coverage     15.28%   28.38%   +13.10%     
   + Complexity      168       53      -115     
   =============================================
     Files          1814     1863       +49     
     Lines         97379    99392     +2013     
     Branches      14893    15146      +253     
   =============================================
   + Hits          14880    28215    +13335     
   + Misses        81367    68427    -12940     
   - Partials       1132     2750     +1618     
   ```
   
   | Flag | Coverage Δ | |
   |---|---|---|
   | integration1 | `26.22% <11.24%> (?)` | |
   | integration2 | `24.77% <0.00%> (?)` | |
   | unittests2 | `?` | |
   
   Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#carryforward-flags-in-the-pull-request-comment) to find out more.
   
   | [Impacted Files](https://codecov.io/gh/apache/pinot/pull/9295?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) | Coverage Δ | |
   |---|---|---|
   | [...gestion/batch/common/BaseSegmentPushJobRunner.java](https://codecov.io/gh/apache/pinot/pull/9295/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtcGx1Z2lucy9waW5vdC1iYXRjaC1pbmdlc3Rpb24vcGlub3QtYmF0Y2gtaW5nZXN0aW9uLWNvbW1vbi9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvcGx1Z2luL2luZ2VzdGlvbi9iYXRjaC9jb21tb24vQmFzZVNlZ21lbnRQdXNoSm9iUnVubmVyLmphdmE=) | `0.00% <0.00%> (ø)` | |
   | [...n/batch/standalone/SegmentGenerationJobRunner.java](https://codecov.io/gh/apache/pinot/pull/9295/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtcGx1Z2lucy9waW5vdC1iYXRjaC1pbmdlc3Rpb24vcGlub3QtYmF0Y2gtaW5nZXN0aW9uLXN0YW5kYWxvbmUvc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL3Bpbm90L3BsdWdpbi9pbmdlc3Rpb24vYmF0Y2gvc3RhbmRhbG9uZS9TZWdtZW50R2VuZXJhdGlvbkpvYlJ1bm5lci5qYXZh) | `0.00% <0.00%> (-71.24%)` | :arrow_down: |
   | [...batch/standalone/SegmentMetadataPushJobRunner.java](https://codecov.io/gh/apache/pinot/pull/9295/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtcGx1Z2lucy9waW5vdC1iYXRjaC1pbmdlc3Rpb24vcGlub3QtYmF0Y2gtaW5nZXN0aW9uLXN0YW5kYWxvbmUvc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL3Bpbm90L3BsdWdpbi9pbmdlc3Rpb24vYmF0Y2gvc3RhbmRhbG9uZS9TZWdtZW50TWV0YWRhdGFQdXNoSm9iUnVubmVyLmphdmE=) | `0.00% <0.00%> (ø)` | |
   | [...tion/batch/standalone/SegmentTarPushJobRunner.java](https://codecov.io/gh/apache/pinot/pull/9295/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtcGx1Z2lucy9waW5vdC1iYXRjaC1pbmdlc3Rpb24vcGlub3QtYmF0Y2gtaW5nZXN0aW9uLXN0YW5kYWxvbmUvc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL3Bpbm90L3BsdWdpbi9pbmdlc3Rpb24vYmF0Y2gvc3RhbmRhbG9uZS9TZWdtZW50VGFyUHVzaEpvYlJ1bm5lci5qYXZh) | `0.00% <0.00%> (ø)` | |
   | [...tion/batch/standalone/SegmentUriPushJobRunner.java](https://codecov.io/gh/apache/pinot/pull/9295/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3QtcGx1Z2lucy9waW5vdC1iYXRjaC1pbmdlc3Rpb24vcGlub3QtYmF0Y2gtaW5nZXN0aW9uLXN0YW5kYWxvbmUvc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL3Bpbm90L3BsdWdpbi9pbmdlc3Rpb24vYmF0Y2gvc3RhbmRhbG9uZS9TZWdtZW50VXJpUHVzaEpvYlJ1bm5lci5qYXZh) | `0.00% <0.00%> (ø)` | |
   | [...t/segment/local/utils/ConsistentDataPushUtils.java](https://codecov.io/gh/apache/pinot/pull/9295/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3Qtc2VnbWVudC1sb2NhbC9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3Qvc2VnbWVudC9sb2NhbC91dGlscy9Db25zaXN0ZW50RGF0YVB1c2hVdGlscy5qYXZh) | `0.00% <0.00%> (ø)` | |
   | [...he/pinot/segment/local/utils/SegmentPushUtils.java](https://codecov.io/gh/apache/pinot/pull/9295/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3Qtc2VnbWVudC1sb2NhbC9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3Qvc2VnbWVudC9sb2NhbC91dGlscy9TZWdtZW50UHVzaFV0aWxzLmphdmE=) | `0.00% <0.00%> (ø)` | |
   | [...i/config/table/ingestion/BatchIngestionConfig.java](https://codecov.io/gh/apache/pinot/pull/9295/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3Qtc3BpL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9waW5vdC9zcGkvY29uZmlnL3RhYmxlL2luZ2VzdGlvbi9CYXRjaEluZ2VzdGlvbkNvbmZpZy5qYXZh) | `0.00% <0.00%> (ø)` | |
   | [...ingestion/batch/spec/SegmentNameGeneratorSpec.java](https://codecov.io/gh/apache/pinot/pull/9295/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3Qtc3BpL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9waW5vdC9zcGkvaW5nZXN0aW9uL2JhdGNoL3NwZWMvU2VnbWVudE5hbWVHZW5lcmF0b3JTcGVjLmphdmE=) | `0.00% <0.00%> (ø)` | |
   | [...g/apache/pinot/spi/utils/IngestionConfigUtils.java](https://codecov.io/gh/apache/pinot/pull/9295/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation#diff-cGlub3Qtc3BpL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9waW5vdC9zcGkvdXRpbHMvSW5nZXN0aW9uQ29uZmlnVXRpbHMuamF2YQ==) | `0.00% <0.00%> (ø)` | |
   | ... and [960 more](https://codecov.io/gh/apache/pinot/pull/9295/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation) | |
   
   :mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=The+Apache+Software+Foundation)
   


-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r959249986


##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java:
##########
@@ -0,0 +1,284 @@
+/**
+ * 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.segment.local.utils;
+
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.common.auth.AuthProviderUtils;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class ConsistentDataPushUtils {
+  private ConsistentDataPushUtils() {
+  }
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentPushUtils.class);
+  private static final FileUploadDownloadClient FILE_UPLOAD_DOWNLOAD_CLIENT = new FileUploadDownloadClient();
+  public static final String SEGMENT_NAME_POSTFIX = "segment.name.postfix";
+
+  /**
+   * Checks for enablement of consistent data push. If enabled, fetch the list of segments to be replaced, then
+   * invoke startReplaceSegments API and returns a map of controller URI to lineage entry IDs.
+   * If not, returns an empty hashmap.
+   */
+  public static Map<URI, String> preUpload(SegmentGenerationJobSpec spec, List<String> segmentsTo)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    boolean consistentDataPushEnabled = consistentDataPushEnabled(spec);
+    LOGGER.info("Consistent data push is: {}", consistentDataPushEnabled ? "enabled" : "disabled");
+    Map<URI, String> uriToLineageEntryIdMap = null;
+    if (consistentDataPushEnabled) {
+      LOGGER.info("Start consistent push for table: " + rawTableName);
+      Map<URI, List<String>> uriToExistingOfflineSegments = getSegmentsToReplace(spec, rawTableName);
+      LOGGER.info("Existing segments for table {}: " + uriToExistingOfflineSegments, rawTableName);
+      LOGGER.info("New segments for table: {}: " + segmentsTo, rawTableName);
+      uriToLineageEntryIdMap = startReplaceSegments(spec, uriToExistingOfflineSegments, segmentsTo);
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * uriToLineageEntryIdMap is non-empty if and only if consistent data push is enabled.
+   * If uriToLineageEntryIdMap is non-empty, end the consistent data push protocol for each controller.
+   */
+  public static void postUpload(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap) {
+    String rawTableName = spec.getTableSpec().getTableName();
+    if (uriToLineageEntryIdMap != null && !uriToLineageEntryIdMap.isEmpty()) {
+      LOGGER.info("End consistent push for table: " + rawTableName);
+      endReplaceSegments(spec, uriToLineageEntryIdMap);
+    }
+  }
+
+  /**
+   * Builds a map of controller URI to startReplaceSegments URI for each Pinot cluster in the spec.
+   */
+  public static Map<URI, URI> getStartReplaceSegmentUris(SegmentGenerationJobSpec spec, String rawTableName) {
+    Map<URI, URI> baseUriToStartReplaceSegmentUriMap = new HashMap<>();
+    for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
+      URI controllerURI;
+      try {
+        controllerURI = new URI(pinotClusterSpec.getControllerURI());
+        baseUriToStartReplaceSegmentUriMap.put(controllerURI,
+            FileUploadDownloadClient.getStartReplaceSegmentsURI(controllerURI, rawTableName,
+                TableType.OFFLINE.toString(), true));
+      } catch (URISyntaxException e) {
+        throw new RuntimeException("Got invalid controller uri - '" + pinotClusterSpec.getControllerURI() + "'");
+      }
+    }
+    return baseUriToStartReplaceSegmentUriMap;
+  }
+
+  /**
+   * Starts consistent data push protocol for each Pinot cluster in the spec.
+   * Returns a map of controller URI to segment lineage entry ID.
+   */
+  public static Map<URI, String> startReplaceSegments(SegmentGenerationJobSpec spec,
+      Map<URI, List<String>> uriToSegmentsFrom, List<String> segmentsTo)
+      throws Exception {
+    Map<URI, String> uriToLineageEntryIdMap = new HashMap<>();
+    String rawTableName = spec.getTableSpec().getTableName();
+    Map<URI, URI> segmentsUris = getStartReplaceSegmentUris(spec, rawTableName);
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    LOGGER.info("Start replace segment URIs: " + segmentsUris);
+
+    int attempts = 1;
+    long retryWaitMs = 1000L;
+
+    for (Map.Entry<URI, URI> entry : segmentsUris.entrySet()) {
+      URI controllerUri = entry.getKey();
+      URI startSegmentUri = entry.getValue();
+      List<String> segmentsFrom = uriToSegmentsFrom.get(controllerUri);
+
+      if (!Collections.disjoint(segmentsFrom, segmentsTo)) {
+        String errorMsg =
+            String.format("Found same segment names when attempting to enable consistent push for table: %s",
+                rawTableName);
+        LOGGER.error("SegmentsFrom: {}", segmentsFrom);
+        LOGGER.error("SegmentsTo: {}", segmentsTo);
+        LOGGER.error(errorMsg);
+        throw new RuntimeException(errorMsg);
+      }
+
+      StartReplaceSegmentsRequest startReplaceSegmentsRequest =
+          new StartReplaceSegmentsRequest(segmentsFrom, segmentsTo);
+      RetryPolicies.exponentialBackoffRetryPolicy(attempts, retryWaitMs, 5).attempt(() -> {
+        try {
+          SimpleHttpResponse response =
+              FILE_UPLOAD_DOWNLOAD_CLIENT.startReplaceSegments(startSegmentUri, startReplaceSegmentsRequest,
+                  authProvider);
+
+          String responseString = response.getResponse();
+          LOGGER.info(
+              "Got response {}: {} while sending start replace segment request for table: {}, uploadURI: {}, request:"
+                  + " {}", response.getStatusCode(), responseString, rawTableName, startSegmentUri,
+              startReplaceSegmentsRequest);
+          String segmentLineageEntryId =
+              JsonUtils.stringToJsonNode(responseString).get("segmentLineageEntryId").asText();
+          uriToLineageEntryIdMap.put(controllerUri, segmentLineageEntryId);
+          return true;
+        } catch (SocketTimeoutException se) {
+          // In case of the timeout, we should re-try.
+          return false;
+        } catch (HttpErrorStatusException e) {
+          if (e.getStatusCode() >= 500) {
+            return false;
+          } else {
+            if (e.getStatusCode() == Response.Status.NOT_FOUND.getStatusCode()) {
+              LOGGER.error("Table: {} not found when sending request: {}", rawTableName, startSegmentUri);
+            }
+            throw e;
+          }
+        }
+      });
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * Ends consistent data push protocol for each Pinot cluster in the spec.
+   */
+  public static void endReplaceSegments(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap) {
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    String rawTableName = spec.getTableSpec().getTableName();
+    for (URI uri : uriToLineageEntryIdMap.keySet()) {
+      String segmentLineageEntryId = uriToLineageEntryIdMap.get(uri);
+      try {
+        FILE_UPLOAD_DOWNLOAD_CLIENT.endReplaceSegments(

Review Comment:
   can we add the retry here as well to keep the consistency with the startReplaceSegments call??



##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java:
##########
@@ -0,0 +1,284 @@
+/**
+ * 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.segment.local.utils;
+
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.common.auth.AuthProviderUtils;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class ConsistentDataPushUtils {

Review Comment:
   I think that we should have the retries for GET API calls to the controller. Can you double-check on how `PinotClusterAccessor:GetRequestCallable` is achieving in our internal push job (pinot-build-push)?



##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java:
##########
@@ -0,0 +1,284 @@
+/**
+ * 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.segment.local.utils;
+
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.common.auth.AuthProviderUtils;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class ConsistentDataPushUtils {
+  private ConsistentDataPushUtils() {
+  }
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentPushUtils.class);
+  private static final FileUploadDownloadClient FILE_UPLOAD_DOWNLOAD_CLIENT = new FileUploadDownloadClient();
+  public static final String SEGMENT_NAME_POSTFIX = "segment.name.postfix";
+
+  /**
+   * Checks for enablement of consistent data push. If enabled, fetch the list of segments to be replaced, then
+   * invoke startReplaceSegments API and returns a map of controller URI to lineage entry IDs.
+   * If not, returns an empty hashmap.
+   */
+  public static Map<URI, String> preUpload(SegmentGenerationJobSpec spec, List<String> segmentsTo)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    boolean consistentDataPushEnabled = consistentDataPushEnabled(spec);
+    LOGGER.info("Consistent data push is: {}", consistentDataPushEnabled ? "enabled" : "disabled");
+    Map<URI, String> uriToLineageEntryIdMap = null;
+    if (consistentDataPushEnabled) {
+      LOGGER.info("Start consistent push for table: " + rawTableName);
+      Map<URI, List<String>> uriToExistingOfflineSegments = getSegmentsToReplace(spec, rawTableName);
+      LOGGER.info("Existing segments for table {}: " + uriToExistingOfflineSegments, rawTableName);
+      LOGGER.info("New segments for table: {}: " + segmentsTo, rawTableName);
+      uriToLineageEntryIdMap = startReplaceSegments(spec, uriToExistingOfflineSegments, segmentsTo);
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * uriToLineageEntryIdMap is non-empty if and only if consistent data push is enabled.
+   * If uriToLineageEntryIdMap is non-empty, end the consistent data push protocol for each controller.
+   */
+  public static void postUpload(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap) {
+    String rawTableName = spec.getTableSpec().getTableName();
+    if (uriToLineageEntryIdMap != null && !uriToLineageEntryIdMap.isEmpty()) {
+      LOGGER.info("End consistent push for table: " + rawTableName);
+      endReplaceSegments(spec, uriToLineageEntryIdMap);
+    }
+  }
+
+  /**
+   * Builds a map of controller URI to startReplaceSegments URI for each Pinot cluster in the spec.
+   */
+  public static Map<URI, URI> getStartReplaceSegmentUris(SegmentGenerationJobSpec spec, String rawTableName) {
+    Map<URI, URI> baseUriToStartReplaceSegmentUriMap = new HashMap<>();
+    for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
+      URI controllerURI;
+      try {
+        controllerURI = new URI(pinotClusterSpec.getControllerURI());
+        baseUriToStartReplaceSegmentUriMap.put(controllerURI,
+            FileUploadDownloadClient.getStartReplaceSegmentsURI(controllerURI, rawTableName,
+                TableType.OFFLINE.toString(), true));
+      } catch (URISyntaxException e) {
+        throw new RuntimeException("Got invalid controller uri - '" + pinotClusterSpec.getControllerURI() + "'");
+      }
+    }
+    return baseUriToStartReplaceSegmentUriMap;
+  }
+
+  /**
+   * Starts consistent data push protocol for each Pinot cluster in the spec.
+   * Returns a map of controller URI to segment lineage entry ID.
+   */
+  public static Map<URI, String> startReplaceSegments(SegmentGenerationJobSpec spec,
+      Map<URI, List<String>> uriToSegmentsFrom, List<String> segmentsTo)
+      throws Exception {
+    Map<URI, String> uriToLineageEntryIdMap = new HashMap<>();
+    String rawTableName = spec.getTableSpec().getTableName();
+    Map<URI, URI> segmentsUris = getStartReplaceSegmentUris(spec, rawTableName);
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    LOGGER.info("Start replace segment URIs: " + segmentsUris);
+
+    int attempts = 1;

Review Comment:
   Can we have a larger number than 1 for a retrial attempt? What's the reasoning behind of doing a retry only for once?



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r960085549


##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java:
##########
@@ -0,0 +1,296 @@
+/**
+ * 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.segment.local.utils;
+
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.common.auth.AuthProviderUtils;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.apache.pinot.spi.utils.retry.RetryPolicy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class ConsistentDataPushUtils {
+  private ConsistentDataPushUtils() {
+  }
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentPushUtils.class);
+  private static final FileUploadDownloadClient FILE_UPLOAD_DOWNLOAD_CLIENT = new FileUploadDownloadClient();
+  private static final RetryPolicy DEFAULT_RETRY_POLICY = RetryPolicies.exponentialBackoffRetryPolicy(5, 10_000L, 2.0);

Review Comment:
   @snleee @jtao15
   Please lmk if these params are appropriate. 



##########
pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java:
##########
@@ -747,6 +754,72 @@ public SimpleHttpResponse uploadSegment(URI uri, String segmentName, InputStream
     return uploadSegment(uri, segmentName, inputStream, null, parameters, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
   }
 
+  /**
+   * Returns a map from a given tableType to a list of segments for that given tableType (OFFLINE or REALTIME)
+   * If tableType is left unspecified, both OFFLINE and REALTIME segments will be returned in the map.
+   */
+  public Map<String, List<String>> getSegments(URI controllerUri, String rawTableName, @Nullable TableType tableType,
+      boolean excludeReplacedSegments) throws Exception {
+    List<String> tableTypes;
+    if (tableType == null) {
+      tableTypes = Arrays.asList(TableType.OFFLINE.toString(), TableType.REALTIME.toString());
+    } else {
+      tableTypes = Arrays.asList(tableType.toString());
+    }
+    ControllerRequestURLBuilder controllerRequestURLBuilder =
+        ControllerRequestURLBuilder.baseUrl(controllerUri.toString());
+    Map<String, List<String>> tableTypeToSegments = new HashMap<>();
+    for (String tableTypeToFilter : tableTypes) {
+      tableTypeToSegments.put(tableTypeToFilter, new ArrayList<>());
+      String uri = controllerRequestURLBuilder.forSegmentListAPIWithTableTypeAndExcludeReplacedSegments(rawTableName,
+          tableTypeToFilter, excludeReplacedSegments);
+      RequestBuilder requestBuilder = RequestBuilder.get(uri).setVersion(HttpVersion.HTTP_1_1);
+      HttpClient.setTimeout(requestBuilder, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
+      RetryPolicies.exponentialBackoffRetryPolicy(5, 10_000L, 2.0).attempt(() -> {

Review Comment:
   Same for this..



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee merged pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
snleee merged PR #9295:
URL: https://github.com/apache/pinot/pull/9295


-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r964268229


##########
pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/BaseSegmentPushJobRunner.java:
##########
@@ -0,0 +1,146 @@
+/**
+ * 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.plugin.ingestion.batch.common;
+
+import java.io.File;
+import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.segment.local.utils.ConsistentDataPushUtils;
+import org.apache.pinot.segment.local.utils.SegmentPushUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.filesystem.PinotFS;
+import org.apache.pinot.spi.filesystem.PinotFSFactory;
+import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+
+public abstract class BaseSegmentPushJobRunner implements IngestionJobRunner {
+
+  protected SegmentGenerationJobSpec _spec;
+  protected String[] _files;
+  protected PinotFS _outputDirFS;
+  protected URI _outputDirURI;
+  protected TableConfig _tableConfig;
+  protected boolean _consistentPushEnabled;
+
+  /**
+   * Initialize BaseSegmentPushJobRunner with SegmentGenerationJobSpec
+   * Checks for required parameters in the spec and enablement of consistent data push.
+   */
+  @Override
+  public void init(SegmentGenerationJobSpec spec) {
+    _spec = spec;
+    if (_spec.getPushJobSpec() == null) {
+      throw new RuntimeException("Missing PushJobSpec");
+    }
+
+    // Read Table spec
+    if (_spec.getTableSpec() == null) {
+      throw new RuntimeException("Missing tableSpec");
+    }
+
+    // Read Table config
+    if (_spec.getTableSpec().getTableConfigURI() == null) {
+      throw new RuntimeException("Missing property 'tableConfigURI' in 'tableSpec'");
+    }
+
+    _tableConfig = SegmentGenerationUtils.getTableConfig(_spec.getTableSpec().getTableConfigURI(), spec.getAuthToken());
+

Review Comment:
   (nit) remove line



##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java:
##########
@@ -371,4 +372,35 @@ private static File generateSegmentMetadataFile(PinotFS fileSystem, URI tarFileU
       FileUtils.deleteQuietly(segmentMetadataDir);
     }
   }
+
+  public static List<String> getSegmentNames(BatchConfigProperties.SegmentPushType pushMode,

Review Comment:
   Why do we need to handle this separately for each `pushMode`? I thought `Map<String, String> segmentsUriToTarPathMap` should be in the same format regardless the pushMode. I thought that we can use the same code and put it to the Base class. Am I missing some detail 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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] jackjlli commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
jackjlli commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r959022101


##########
pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java:
##########
@@ -747,6 +751,54 @@ public SimpleHttpResponse uploadSegment(URI uri, String segmentName, InputStream
     return uploadSegment(uri, segmentName, inputStream, null, parameters, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
   }
 
+  /**
+   * Returns a map from a given tableType to a list of segments for that given tableType (OFFLINE or REALTIME)
+   * If tableType is left unspecified, both OFFLINE and REALTIME segments will be returned in the map.
+   */
+  public Map<String, List<String>> getSegments(URI uri, String rawTableName, @Nullable String tableType,
+      boolean excludeReplacedSegments)
+      throws URISyntaxException, IOException {
+    List<String> tableTypes;
+    if (tableType == null || tableType.isEmpty()) {
+      tableTypes = Arrays.asList(TableType.OFFLINE.toString(), TableType.REALTIME.toString());
+    } else {
+      tableTypes = Arrays.asList(tableType);

Review Comment:
   should we validate the correctness of the table type here if it's non-empty?



##########
pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java:
##########
@@ -747,6 +751,54 @@ public SimpleHttpResponse uploadSegment(URI uri, String segmentName, InputStream
     return uploadSegment(uri, segmentName, inputStream, null, parameters, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
   }
 
+  /**
+   * Returns a map from a given tableType to a list of segments for that given tableType (OFFLINE or REALTIME)
+   * If tableType is left unspecified, both OFFLINE and REALTIME segments will be returned in the map.
+   */
+  public Map<String, List<String>> getSegments(URI uri, String rawTableName, @Nullable String tableType,
+      boolean excludeReplacedSegments)
+      throws URISyntaxException, IOException {
+    List<String> tableTypes;
+    if (tableType == null || tableType.isEmpty()) {
+      tableTypes = Arrays.asList(TableType.OFFLINE.toString(), TableType.REALTIME.toString());
+    } else {
+      tableTypes = Arrays.asList(tableType);
+    }
+    ControllerRequestURLBuilder controllerRequestURLBuilder = ControllerRequestURLBuilder.baseUrl(uri.toString());
+    Map<String, List<String>> tableTypeToSegments = new HashMap<>();
+    for (String tableTypeToFilter : tableTypes) {
+      List<String> segments = new ArrayList<>();
+      RequestBuilder requestBuilder = RequestBuilder.get(
+          controllerRequestURLBuilder.forSegmentListAPIWithTableTypeAndExcludeReplacedSegments(rawTableName,
+              tableTypeToFilter, excludeReplacedSegments)).setVersion(HttpVersion.HTTP_1_1);
+      HttpClient.setTimeout(requestBuilder, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
+      SimpleHttpResponse response;
+      try {
+        response = HttpClient.wrapAndThrowHttpException(_httpClient.sendRequest(requestBuilder.build()));
+      } catch (HttpErrorStatusException e) {
+        tableTypeToSegments.put(tableTypeToFilter, segments);
+        continue;
+      }
+      String responseString = response.getResponse();
+      JsonNode responseJsonNode = JsonUtils.stringToJsonNode(responseString);
+      Iterator<JsonNode> responseElements = responseJsonNode.elements();
+      while (responseElements.hasNext()) {

Review Comment:
   Can we encapsulate the logic of fetching the segment names into a method, or is there any existing logic doing that right now?



##########
pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java:
##########
@@ -747,6 +751,54 @@ public SimpleHttpResponse uploadSegment(URI uri, String segmentName, InputStream
     return uploadSegment(uri, segmentName, inputStream, null, parameters, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
   }
 
+  /**
+   * Returns a map from a given tableType to a list of segments for that given tableType (OFFLINE or REALTIME)
+   * If tableType is left unspecified, both OFFLINE and REALTIME segments will be returned in the map.
+   */
+  public Map<String, List<String>> getSegments(URI uri, String rawTableName, @Nullable String tableType,
+      boolean excludeReplacedSegments)
+      throws URISyntaxException, IOException {
+    List<String> tableTypes;
+    if (tableType == null || tableType.isEmpty()) {
+      tableTypes = Arrays.asList(TableType.OFFLINE.toString(), TableType.REALTIME.toString());
+    } else {
+      tableTypes = Arrays.asList(tableType);
+    }
+    ControllerRequestURLBuilder controllerRequestURLBuilder = ControllerRequestURLBuilder.baseUrl(uri.toString());
+    Map<String, List<String>> tableTypeToSegments = new HashMap<>();
+    for (String tableTypeToFilter : tableTypes) {
+      List<String> segments = new ArrayList<>();
+      RequestBuilder requestBuilder = RequestBuilder.get(
+          controllerRequestURLBuilder.forSegmentListAPIWithTableTypeAndExcludeReplacedSegments(rawTableName,
+              tableTypeToFilter, excludeReplacedSegments)).setVersion(HttpVersion.HTTP_1_1);
+      HttpClient.setTimeout(requestBuilder, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
+      SimpleHttpResponse response;
+      try {
+        response = HttpClient.wrapAndThrowHttpException(_httpClient.sendRequest(requestBuilder.build()));
+      } catch (HttpErrorStatusException e) {
+        tableTypeToSegments.put(tableTypeToFilter, segments);

Review Comment:
   It'd be good to print out the exception message here instead of swallowing it silently.



##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java:
##########
@@ -371,4 +372,36 @@ private static File generateSegmentMetadataFile(PinotFS fileSystem, URI tarFileU
       FileUtils.deleteQuietly(segmentMetadataDir);
     }
   }
+
+  public static List<String> getSegmentNames(BatchConfigProperties.SegmentPushType pushMode,
+      Map<String, String> segmentsUriToTarPathMap) {
+    List<String> segmentNames = new ArrayList<>();

Review Comment:
   nit: you can instantiate the array list with the size of the `segmentsUriToTarPathMap` map, so that you don't have to worry about the resizing.



##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java:
##########
@@ -371,4 +372,36 @@ private static File generateSegmentMetadataFile(PinotFS fileSystem, URI tarFileU
       FileUtils.deleteQuietly(segmentMetadataDir);
     }
   }
+
+  public static List<String> getSegmentNames(BatchConfigProperties.SegmentPushType pushMode,
+      Map<String, String> segmentsUriToTarPathMap) {
+    List<String> segmentNames = new ArrayList<>();
+    if (pushMode.equals(BatchConfigProperties.SegmentPushType.TAR) || pushMode.equals(
+        BatchConfigProperties.SegmentPushType.METADATA)) {
+      for (String tarFilePath : segmentsUriToTarPathMap.values()) {
+        File tarFile = new File(tarFilePath);
+        String fileName = tarFile.getName();
+        Preconditions.checkArgument(fileName.endsWith(Constants.TAR_GZ_FILE_EXT));
+        String segmentName = getSegmentNameFromPath(fileName);
+        segmentNames.add(segmentName);
+      }
+    }
+    if (pushMode.equals(BatchConfigProperties.SegmentPushType.URI)) {
+      for (String segmentUri : segmentsUriToTarPathMap.keySet()) {
+        Preconditions.checkArgument(segmentUri.endsWith(Constants.TAR_GZ_FILE_EXT));
+        String segmentName = getSegmentNameFromPath(segmentUri);
+        segmentNames.add(segmentName);
+      }
+    }
+    return segmentNames;
+  }
+
+  /**
+   * Obtain segment name given filePath by reading from after the last slash (if present) up to and before the tar
+   * extension.
+   */
+  public static String getSegmentNameFromPath(String filePath) {
+    int startIndex = filePath.contains("/") ? filePath.lastIndexOf("/") + 1 : 0;

Review Comment:
   nit: the `lastIndexOf()` method will return -1 if it doesn't exist, so you don't have to check whether the string `contains` "/" in the first place.



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] agavra commented on pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
agavra commented on PR #9295:
URL: https://github.com/apache/pinot/pull/9295#issuecomment-1264885127

   +1 to @Jackie-Jiang, I think this commit breaks the quickstart for batch import: https://docs.pinot.apache.org/basics/getting-started/pushing-your-data-to-pinot


-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r959048943


##########
pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java:
##########
@@ -747,6 +751,54 @@ public SimpleHttpResponse uploadSegment(URI uri, String segmentName, InputStream
     return uploadSegment(uri, segmentName, inputStream, null, parameters, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
   }
 
+  /**
+   * Returns a map from a given tableType to a list of segments for that given tableType (OFFLINE or REALTIME)
+   * If tableType is left unspecified, both OFFLINE and REALTIME segments will be returned in the map.
+   */
+  public Map<String, List<String>> getSegments(URI uri, String rawTableName, @Nullable String tableType,
+      boolean excludeReplacedSegments)
+      throws URISyntaxException, IOException {
+    List<String> tableTypes;
+    if (tableType == null || tableType.isEmpty()) {
+      tableTypes = Arrays.asList(TableType.OFFLINE.toString(), TableType.REALTIME.toString());
+    } else {
+      tableTypes = Arrays.asList(tableType);

Review Comment:
   Refactored the code to take in `tableType` as `TableType` rather than `String`. That should prevent illegal table types. 



##########
pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java:
##########
@@ -747,6 +751,54 @@ public SimpleHttpResponse uploadSegment(URI uri, String segmentName, InputStream
     return uploadSegment(uri, segmentName, inputStream, null, parameters, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
   }
 
+  /**
+   * Returns a map from a given tableType to a list of segments for that given tableType (OFFLINE or REALTIME)
+   * If tableType is left unspecified, both OFFLINE and REALTIME segments will be returned in the map.
+   */
+  public Map<String, List<String>> getSegments(URI uri, String rawTableName, @Nullable String tableType,
+      boolean excludeReplacedSegments)
+      throws URISyntaxException, IOException {
+    List<String> tableTypes;
+    if (tableType == null || tableType.isEmpty()) {
+      tableTypes = Arrays.asList(TableType.OFFLINE.toString(), TableType.REALTIME.toString());
+    } else {
+      tableTypes = Arrays.asList(tableType);

Review Comment:
   Refactored the code to take in `tableType` as `TableType` rather than `String`. That should prevent illegal table types, thanks. 



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r958990071


##########
pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/BaseSegmentPushJobRunner.java:
##########
@@ -0,0 +1,142 @@
+/**
+ * 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.plugin.ingestion.batch.common;
+
+import java.io.File;
+import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.segment.local.utils.ConsistentDataPushUtils;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.filesystem.PinotFS;
+import org.apache.pinot.spi.filesystem.PinotFSFactory;
+import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+
+public abstract class BaseSegmentPushJobRunner implements IngestionJobRunner {
+
+  protected SegmentGenerationJobSpec _spec;
+  protected String[] _files;
+  protected PinotFS _outputDirFS;
+  protected URI _outputDirURI;
+  protected List<String> _segmentsToPush = new ArrayList<>();
+  protected Map<String, String> _segmentUriToTarPathMap;
+  protected boolean _consistentPushEnabled;
+
+  /**
+   * Initialize BaseSegmentPushJobRunner with SegmentGenerationJobSpec
+   * Checks for required parameters in the spec and enablement of consistent data push.
+   */
+  @Override
+  public void init(SegmentGenerationJobSpec spec) {
+    _spec = spec;
+    if (_spec.getPushJobSpec() == null) {
+      throw new RuntimeException("Missing PushJobSpec");
+    }
+
+    // Read Table spec
+    if (_spec.getTableSpec() == null) {
+      throw new RuntimeException("Missing tableSpec");
+    }
+
+    // Read Table config
+    if (_spec.getTableSpec().getTableConfigURI() == null) {
+      throw new RuntimeException("Missing property 'tableConfigURI' in 'tableSpec'");
+    }
+
+    _consistentPushEnabled = ConsistentDataPushUtils.consistentDataPushEnabled(_spec);
+  }
+
+  /**
+   * Initialize filesystems and obtain the raw input files for upload.
+   */
+  public void initFileSys() {
+    // init all file systems
+    List<PinotFSSpec> pinotFSSpecs = _spec.getPinotFSSpecs();
+    for (PinotFSSpec pinotFSSpec : pinotFSSpecs) {
+      PinotFSFactory.register(pinotFSSpec.getScheme(), pinotFSSpec.getClassName(), new PinotConfiguration(pinotFSSpec));
+    }
+
+    // Get outputFS for writing output Pinot segments
+    try {
+      _outputDirURI = new URI(_spec.getOutputDirURI());
+      if (_outputDirURI.getScheme() == null) {
+        _outputDirURI = new File(_spec.getOutputDirURI()).toURI();
+      }
+    } catch (URISyntaxException e) {
+      throw new RuntimeException("outputDirURI is not valid - '" + _spec.getOutputDirURI() + "'");
+    }
+    _outputDirFS = PinotFSFactory.create(_outputDirURI.getScheme());
+
+    // Get list of files to process
+    try {
+      _files = _outputDirFS.listFiles(_outputDirURI, true);
+    } catch (IOException e) {
+      throw new RuntimeException("Unable to list all files under outputDirURI - '" + _outputDirURI + "'");
+    }
+  }
+
+  /**
+   * Populates either _segmentsToPush or _segmentUriToTarPathMap based on push type.
+   */
+  public abstract void getSegmentsToPush();
+
+  /**
+   * Returns segmentsTo based on segments obtained from getSegmentsToPush.
+   * The result will to be supplied to the segment replacement protocol when consistent data push is enabled.
+   */
+  public abstract List<String> getSegmentsTo();
+
+  /**
+   * Upload segment obtained by getSegmentsToPush.
+   */
+  public abstract void uploadSegments()
+      throws Exception;
+
+  /**
+   * Runs the main logic of the segment push job runner.
+   * First initialize the filesystem, then upload the segments, while optionally configured to be wrapped around by
+   * the consistent data push protocol.
+   */
+  @Override
+  public void run() {
+    initFileSys();
+    Map<URI, String> uriToLineageEntryIdMap = null;
+    try {
+      getSegmentsToPush();

Review Comment:
   Thanks, this is a great suggestion that I missed. Addressed as recommended!



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r964080686


##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java:
##########
@@ -0,0 +1,285 @@
+/**
+ * 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.segment.local.utils;
+
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.common.auth.AuthProviderUtils;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.apache.pinot.spi.utils.retry.RetryPolicy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class ConsistentDataPushUtils {
+  private ConsistentDataPushUtils() {
+  }
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentPushUtils.class);
+  private static final FileUploadDownloadClient FILE_UPLOAD_DOWNLOAD_CLIENT = new FileUploadDownloadClient();
+  private static final RetryPolicy DEFAULT_RETRY_POLICY = RetryPolicies.exponentialBackoffRetryPolicy(5, 10_000L, 2.0);
+  public static final String SEGMENT_NAME_POSTFIX = "segment.name.postfix";
+
+  /**
+   * Checks for enablement of consistent data push. If enabled, fetch the list of segments to be replaced, then
+   * invoke startReplaceSegments API and returns a map of controller URI to lineage entry IDs.
+   * If not, returns an empty hashmap.
+   */
+  public static Map<URI, String> preUpload(SegmentGenerationJobSpec spec, List<String> segmentsTo)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    boolean consistentDataPushEnabled = consistentDataPushEnabled(spec);
+    LOGGER.info("Consistent data push is: {}", consistentDataPushEnabled ? "enabled" : "disabled");
+    Map<URI, String> uriToLineageEntryIdMap = null;
+    if (consistentDataPushEnabled) {
+      LOGGER.info("Start consistent push for table: " + rawTableName);
+      Map<URI, List<String>> uriToExistingOfflineSegments = getSegmentsToReplace(spec, rawTableName);
+      LOGGER.info("Existing segments for table {}: " + uriToExistingOfflineSegments, rawTableName);
+      LOGGER.info("New segments for table: {}: " + segmentsTo, rawTableName);
+      uriToLineageEntryIdMap = startReplaceSegments(spec, uriToExistingOfflineSegments, segmentsTo);
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * uriToLineageEntryIdMap is non-empty if and only if consistent data push is enabled.
+   * If uriToLineageEntryIdMap is non-empty, end the consistent data push protocol for each controller.
+   */
+  public static void postUpload(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    if (uriToLineageEntryIdMap != null && !uriToLineageEntryIdMap.isEmpty()) {
+      LOGGER.info("End consistent push for table: " + rawTableName);
+      endReplaceSegments(spec, uriToLineageEntryIdMap);
+    }
+  }
+
+  /**
+   * Builds a map of controller URI to startReplaceSegments URI for each Pinot cluster in the spec.
+   */
+  public static Map<URI, URI> getStartReplaceSegmentUris(SegmentGenerationJobSpec spec, String rawTableName) {
+    Map<URI, URI> baseUriToStartReplaceSegmentUriMap = new HashMap<>();
+    for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
+      URI controllerURI;
+      try {
+        controllerURI = new URI(pinotClusterSpec.getControllerURI());
+        baseUriToStartReplaceSegmentUriMap.put(controllerURI,
+            FileUploadDownloadClient.getStartReplaceSegmentsURI(controllerURI, rawTableName,
+                TableType.OFFLINE.toString(), true));
+      } catch (URISyntaxException e) {
+        throw new RuntimeException("Got invalid controller uri - '" + pinotClusterSpec.getControllerURI() + "'");
+      }
+    }
+    return baseUriToStartReplaceSegmentUriMap;
+  }
+
+  /**
+   * Starts consistent data push protocol for each Pinot cluster in the spec.
+   * Returns a map of controller URI to segment lineage entry ID.
+   */
+  public static Map<URI, String> startReplaceSegments(SegmentGenerationJobSpec spec,
+      Map<URI, List<String>> uriToSegmentsFrom, List<String> segmentsTo)
+      throws Exception {
+    Map<URI, String> uriToLineageEntryIdMap = new HashMap<>();
+    String rawTableName = spec.getTableSpec().getTableName();
+    Map<URI, URI> segmentsUris = getStartReplaceSegmentUris(spec, rawTableName);
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    LOGGER.info("Start replace segment URIs: " + segmentsUris);
+
+    for (Map.Entry<URI, URI> entry : segmentsUris.entrySet()) {
+      URI controllerUri = entry.getKey();
+      URI startSegmentUri = entry.getValue();
+      List<String> segmentsFrom = uriToSegmentsFrom.get(controllerUri);
+
+      StartReplaceSegmentsRequest startReplaceSegmentsRequest =
+          new StartReplaceSegmentsRequest(segmentsFrom, segmentsTo);
+      DEFAULT_RETRY_POLICY.attempt(() -> {
+        try {
+          SimpleHttpResponse response =
+              FILE_UPLOAD_DOWNLOAD_CLIENT.startReplaceSegments(startSegmentUri, startReplaceSegmentsRequest,
+                  authProvider);
+          String responseString = response.getResponse();
+          LOGGER.info(
+              "Got response {}: {} while sending start replace segment request for table: {}, uploadURI: {}, request:"
+                  + " {}", response.getStatusCode(), responseString, rawTableName, startSegmentUri,
+              startReplaceSegmentsRequest);
+          String segmentLineageEntryId =
+              JsonUtils.stringToJsonNode(responseString).get("segmentLineageEntryId").asText();
+          uriToLineageEntryIdMap.put(controllerUri, segmentLineageEntryId);
+          return true;
+        } catch (SocketTimeoutException se) {
+          // In case of the timeout, we should re-try.
+          return false;
+        } catch (HttpErrorStatusException e) {
+          if (e.getStatusCode() >= 500) {
+            return false;
+          } else {
+            if (e.getStatusCode() == Response.Status.NOT_FOUND.getStatusCode()) {
+              LOGGER.error("Table: {} not found when sending request: {}", rawTableName, startSegmentUri);
+            }
+            throw e;
+          }
+        }
+      });
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * Ends consistent data push protocol for each Pinot cluster in the spec.
+   */
+  public static void endReplaceSegments(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap)
+      throws Exception {
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    String rawTableName = spec.getTableSpec().getTableName();
+    for (URI controllerUri : uriToLineageEntryIdMap.keySet()) {
+      String segmentLineageEntryId = uriToLineageEntryIdMap.get(controllerUri);
+      URI uri =
+          FileUploadDownloadClient.getEndReplaceSegmentsURI(controllerUri, rawTableName, TableType.OFFLINE.toString(),
+              segmentLineageEntryId);
+      DEFAULT_RETRY_POLICY.attempt(() -> {
+        try {
+          SimpleHttpResponse response =
+              FILE_UPLOAD_DOWNLOAD_CLIENT.endReplaceSegments(uri, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS, authProvider);
+          LOGGER.info("Got response {}: {} while sending end replace segment request for table: {}, uploadURI: {}",
+              response.getStatusCode(), response.getResponse(), rawTableName, uri);
+          return true;
+        } catch (SocketTimeoutException se) {
+          // In case of the timeout, we should re-try.
+          return false;
+        } catch (HttpErrorStatusException e) {
+          if (e.getStatusCode() >= 500) {
+            return false;
+          } else {
+            throw e;
+          }
+        }
+      });
+    }
+  }
+
+  /**
+   * Revert segment lineage entry when exception gets caught. This revert request is called at best effort.
+   * If the revert call fails at this point, the next startReplaceSegment call will do the cleanup
+   * by marking the previous entry to "REVERTED" and cleaning up the leftover segments.
+   */
+  public static void handleUploadException(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap,
+      Exception exception) {
+    if (uriToLineageEntryIdMap != null) {
+      LOGGER.error("Exception when pushing segments. Marking segment lineage entry to 'REVERTED'.", exception);
+      String rawTableName = spec.getTableSpec().getTableName();
+      for (Map.Entry<URI, String> entry : uriToLineageEntryIdMap.entrySet()) {
+        String segmentLineageEntryId = entry.getValue();
+        try {
+          URI uri = FileUploadDownloadClient.getRevertReplaceSegmentsURI(entry.getKey(), rawTableName,
+              TableType.OFFLINE.name(), segmentLineageEntryId, true);
+          SimpleHttpResponse response = FILE_UPLOAD_DOWNLOAD_CLIENT.revertReplaceSegments(uri);
+          LOGGER.info("Got response {}: {} while sending revert replace segment request for table: {}, uploadURI: {}",
+              response.getStatusCode(), response.getResponse(), rawTableName, entry.getKey());
+        } catch (URISyntaxException | HttpErrorStatusException | IOException e) {
+          LOGGER.error("Exception when sending revert replace segment request to controller: {} for table: {}",
+              entry.getKey(), rawTableName, e);
+        }
+      }
+    }
+  }
+
+  public static TableConfig getTableConfig(SegmentGenerationJobSpec spec) {
+    String rawTableName = spec.getTableSpec().getTableName();
+    TableConfig tableConfig =
+        SegmentGenerationUtils.getTableConfig(spec.getTableSpec().getTableConfigURI(), spec.getAuthToken());
+    Preconditions.checkState(tableConfig != null, "Failed to find table config for table: %s", rawTableName);
+    return tableConfig;
+  }
+
+  public static boolean consistentDataPushEnabled(SegmentGenerationJobSpec spec) {
+    TableConfig tableConfig = getTableConfig(spec);
+    // Enable consistent data push only if "consistentDataPush" is set to true in batch ingestion config and the
+    // table is REFRESH use case.
+    // TODO: Remove the check for REFRESH when we support consistent push for APPEND table
+    return "REFRESH".equalsIgnoreCase(IngestionConfigUtils.getBatchSegmentIngestionType(tableConfig))
+        && IngestionConfigUtils.getBatchSegmentIngestionConsistentDataPushEnabled(tableConfig);
+  }
+
+  /**
+   * Returns a map of controller URI to a list of existing OFFLINE segments.
+   */
+  public static Map<URI, List<String>> getSegmentsToReplace(SegmentGenerationJobSpec spec, String rawTableName) {
+    Map<URI, List<String>> uriToOfflineSegments = new HashMap<>();
+    for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
+      URI controllerURI;
+      List<String> offlineSegments;
+      try {
+        controllerURI = new URI(pinotClusterSpec.getControllerURI());
+        Map<String, List<String>> segments =
+            FILE_UPLOAD_DOWNLOAD_CLIENT.getSegments(controllerURI, rawTableName, TableType.OFFLINE, true);
+        offlineSegments = segments.get(TableType.OFFLINE.toString());
+        uriToOfflineSegments.put(controllerURI, offlineSegments);
+      } catch (URISyntaxException e) {
+        throw new RuntimeException("Got invalid controller uri - '" + pinotClusterSpec.getControllerURI() + "'");
+      } catch (Exception e) {
+        e.printStackTrace();
+      }
+    }
+    return uriToOfflineSegments;
+  }
+
+  /**
+   * If consistent data push is enabled, append current timestamp to existing configured segment name postfix, if
+   * configured, to make segment name unique.
+   */
+  public static void configureSegmentPostfix(SegmentGenerationJobSpec spec) {
+    SegmentNameGeneratorSpec segmentNameGeneratorSpec = spec.getSegmentNameGeneratorSpec();
+    if (consistentDataPushEnabled(spec)) {

Review Comment:
   Good catch, reflected. 



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r958990701


##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentUploadIntegrationTest.java:
##########
@@ -188,6 +202,133 @@ public void testUploadAndQuery()
     testCountStar(numDocs);
   }
 
+  /**
+   * Runs both SegmentMetadataPushJobRunner and SegmentTarPushJobRunner while enabling consistent data push.
+   * Checks that segments are properly loaded and segment lineage entry were also in expected states.
+   */
+  @Test
+  public void testUploadAndQueryWithConsistentPush()
+      throws Exception {
+    // Create and upload the schema and table config
+    Schema schema = createSchema();
+    addSchema(schema);
+    TableConfig offlineTableConfig = createOfflineTableConfigWithConsistentPush();
+    addTableConfig(offlineTableConfig);
+
+    List<File> avroFiles = getAllAvroFiles();
+
+    ClusterIntegrationTestUtils.buildSegmentFromAvro(avroFiles.get(0), offlineTableConfig, schema, "_with_move",
+        _segmentDir, _tarDir);
+
+    // First test standalone metadata push job runner
+    BaseSegmentPushJobRunner runner = new SegmentMetadataPushJobRunner();
+    SegmentGenerationJobSpec jobSpec = new SegmentGenerationJobSpec();
+    PushJobSpec pushJobSpec = new PushJobSpec();
+    pushJobSpec.setCopyToDeepStoreForMetadataPush(true);
+    jobSpec.setPushJobSpec(pushJobSpec);
+    PinotFSSpec fsSpec = new PinotFSSpec();
+    fsSpec.setScheme("file");
+    fsSpec.setClassName("org.apache.pinot.spi.filesystem.LocalPinotFS");
+    jobSpec.setPinotFSSpecs(Lists.newArrayList(fsSpec));
+    jobSpec.setOutputDirURI(_tarDir.getAbsolutePath());
+    TableSpec tableSpec = new TableSpec();
+    tableSpec.setTableName(DEFAULT_TABLE_NAME);
+    tableSpec.setTableConfigURI(_controllerRequestURLBuilder.forUpdateTableConfig(DEFAULT_TABLE_NAME));
+    jobSpec.setTableSpec(tableSpec);
+    PinotClusterSpec clusterSpec = new PinotClusterSpec();
+    clusterSpec.setControllerURI(_controllerBaseApiUrl);
+    jobSpec.setPinotClusterSpecs(new PinotClusterSpec[]{clusterSpec});
+
+    File dataDir = new File(_controllerConfig.getDataDir());
+    File dataDirSegments = new File(dataDir, DEFAULT_TABLE_NAME);
+
+    Assert.assertEquals(_tarDir.listFiles().length, 1);
+
+    runner.init(jobSpec);
+    runner.run();
+
+    // Segment should be seen in dataDir
+    Assert.assertTrue(dataDirSegments.exists());
+    Assert.assertEquals(dataDirSegments.listFiles().length, 1);
+    Assert.assertEquals(_tarDir.listFiles().length, 1);
+
+    // test segment loaded
+    JsonNode segmentsList = getSegmentsList();
+    Assert.assertEquals(segmentsList.size(), 1);
+    String segmentNameWithMove = segmentsList.get(0).asText();

Review Comment:
   done



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r957981334


##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentUploadIntegrationTest.java:
##########
@@ -81,10 +89,15 @@ protected List<String> getBloomFilterColumns() {
     return null;
   }
 
+  @BeforeMethod
+  public void setUpTest()

Review Comment:
   I think that this is not actually necessary.
   
   We can call `TestUtils.ensureDirectoriesExistAndEmpty()` function in `setUp()` and clean up those in `cleanUp()`
   
   For 2 different tests, we can use different table names.



##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentUploadIntegrationTest.java:
##########
@@ -217,10 +358,15 @@ public Boolean apply(@Nullable Void aVoid) {
     }, 100L, 300_000, "Failed to load " + countStarResult + " documents", true);
   }
 
+  @AfterMethod

Review Comment:
   This is also not necessary if we use different table names.



##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentUploadIntegrationTest.java:
##########
@@ -188,6 +202,133 @@ public void testUploadAndQuery()
     testCountStar(numDocs);
   }
 
+  /**
+   * Runs both SegmentMetadataPushJobRunner and SegmentTarPushJobRunner while enabling consistent data push.
+   * Checks that segments are properly loaded and segment lineage entry were also in expected states.
+   */
+  @Test
+  public void testUploadAndQueryWithConsistentPush()
+      throws Exception {
+    // Create and upload the schema and table config
+    Schema schema = createSchema();
+    addSchema(schema);
+    TableConfig offlineTableConfig = createOfflineTableConfigWithConsistentPush();
+    addTableConfig(offlineTableConfig);
+
+    List<File> avroFiles = getAllAvroFiles();
+
+    ClusterIntegrationTestUtils.buildSegmentFromAvro(avroFiles.get(0), offlineTableConfig, schema, "_with_move",
+        _segmentDir, _tarDir);
+
+    // First test standalone metadata push job runner
+    BaseSegmentPushJobRunner runner = new SegmentMetadataPushJobRunner();
+    SegmentGenerationJobSpec jobSpec = new SegmentGenerationJobSpec();
+    PushJobSpec pushJobSpec = new PushJobSpec();
+    pushJobSpec.setCopyToDeepStoreForMetadataPush(true);
+    jobSpec.setPushJobSpec(pushJobSpec);
+    PinotFSSpec fsSpec = new PinotFSSpec();
+    fsSpec.setScheme("file");
+    fsSpec.setClassName("org.apache.pinot.spi.filesystem.LocalPinotFS");
+    jobSpec.setPinotFSSpecs(Lists.newArrayList(fsSpec));
+    jobSpec.setOutputDirURI(_tarDir.getAbsolutePath());
+    TableSpec tableSpec = new TableSpec();
+    tableSpec.setTableName(DEFAULT_TABLE_NAME);
+    tableSpec.setTableConfigURI(_controllerRequestURLBuilder.forUpdateTableConfig(DEFAULT_TABLE_NAME));
+    jobSpec.setTableSpec(tableSpec);
+    PinotClusterSpec clusterSpec = new PinotClusterSpec();
+    clusterSpec.setControllerURI(_controllerBaseApiUrl);
+    jobSpec.setPinotClusterSpecs(new PinotClusterSpec[]{clusterSpec});
+
+    File dataDir = new File(_controllerConfig.getDataDir());
+    File dataDirSegments = new File(dataDir, DEFAULT_TABLE_NAME);
+
+    Assert.assertEquals(_tarDir.listFiles().length, 1);
+
+    runner.init(jobSpec);
+    runner.run();
+
+    // Segment should be seen in dataDir
+    Assert.assertTrue(dataDirSegments.exists());
+    Assert.assertEquals(dataDirSegments.listFiles().length, 1);
+    Assert.assertEquals(_tarDir.listFiles().length, 1);
+
+    // test segment loaded
+    JsonNode segmentsList = getSegmentsList();
+    Assert.assertEquals(segmentsList.size(), 1);
+    String segmentNameWithMove = segmentsList.get(0).asText();
+    Assert.assertTrue(segmentNameWithMove.endsWith("_with_move"));
+    long numDocs = getNumDocs(segmentNameWithMove);
+    testCountStar(numDocs);
+
+    // Fetch segment lineage entry after running segment metadata push with consistent push enabled.
+    String segmentLineageResponse = ControllerTest.sendGetRequest(
+        ControllerRequestURLBuilder.baseUrl(_controllerBaseApiUrl)
+            .forListAllSegmentLineages(DEFAULT_TABLE_NAME, TableType.OFFLINE.toString()));
+    // Segment lineage should be in completed state.
+    Assert.assertTrue(segmentLineageResponse.contains("\"state\":\"COMPLETED\""));
+    // SegmentsFrom should be empty as we started with a blank table.
+    Assert.assertTrue(segmentLineageResponse.contains("\"segmentsFrom\":[]"));
+    // SegmentsTo should contain uploaded segment.
+    Assert.assertTrue(segmentLineageResponse.contains("\"segmentsTo\":[\"" + segmentNameWithMove + "\"]"));
+
+    // Clear segment and tar dir
+    for (File segment : _segmentDir.listFiles()) {
+      FileUtils.deleteQuietly(segment);
+    }
+    for (File tar : _tarDir.listFiles()) {
+      FileUtils.deleteQuietly(tar);
+    }
+
+    ClusterIntegrationTestUtils.buildSegmentFromAvro(avroFiles.get(1), offlineTableConfig, schema, "_without_move",
+        _segmentDir, _tarDir);
+    jobSpec.setPushJobSpec(new PushJobSpec());
+
+    // Now test standalone tar push job runner
+    runner = new SegmentTarPushJobRunner();
+
+    Assert.assertEquals(dataDirSegments.listFiles().length, 1);
+    Assert.assertEquals(_tarDir.listFiles().length, 1);
+
+    runner.init(jobSpec);
+    runner.run();
+
+    // should not see new segments in dataDir
+    //Assert.assertEquals(dataDirSegments.listFiles().length, 1);

Review Comment:
   delete?



##########
pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/BaseSegmentPushJobRunner.java:
##########
@@ -0,0 +1,142 @@
+/**
+ * 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.plugin.ingestion.batch.common;
+
+import java.io.File;
+import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.segment.local.utils.ConsistentDataPushUtils;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.filesystem.PinotFS;
+import org.apache.pinot.spi.filesystem.PinotFSFactory;
+import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+
+public abstract class BaseSegmentPushJobRunner implements IngestionJobRunner {
+
+  protected SegmentGenerationJobSpec _spec;
+  protected String[] _files;
+  protected PinotFS _outputDirFS;
+  protected URI _outputDirURI;
+  protected List<String> _segmentsToPush = new ArrayList<>();
+  protected Map<String, String> _segmentUriToTarPathMap;
+  protected boolean _consistentPushEnabled;
+
+  /**
+   * Initialize BaseSegmentPushJobRunner with SegmentGenerationJobSpec
+   * Checks for required parameters in the spec and enablement of consistent data push.
+   */
+  @Override
+  public void init(SegmentGenerationJobSpec spec) {
+    _spec = spec;
+    if (_spec.getPushJobSpec() == null) {
+      throw new RuntimeException("Missing PushJobSpec");
+    }
+
+    // Read Table spec
+    if (_spec.getTableSpec() == null) {
+      throw new RuntimeException("Missing tableSpec");
+    }
+
+    // Read Table config
+    if (_spec.getTableSpec().getTableConfigURI() == null) {
+      throw new RuntimeException("Missing property 'tableConfigURI' in 'tableSpec'");
+    }
+
+    _consistentPushEnabled = ConsistentDataPushUtils.consistentDataPushEnabled(_spec);
+  }
+
+  /**
+   * Initialize filesystems and obtain the raw input files for upload.
+   */
+  public void initFileSys() {
+    // init all file systems
+    List<PinotFSSpec> pinotFSSpecs = _spec.getPinotFSSpecs();
+    for (PinotFSSpec pinotFSSpec : pinotFSSpecs) {
+      PinotFSFactory.register(pinotFSSpec.getScheme(), pinotFSSpec.getClassName(), new PinotConfiguration(pinotFSSpec));
+    }
+
+    // Get outputFS for writing output Pinot segments
+    try {
+      _outputDirURI = new URI(_spec.getOutputDirURI());
+      if (_outputDirURI.getScheme() == null) {
+        _outputDirURI = new File(_spec.getOutputDirURI()).toURI();
+      }
+    } catch (URISyntaxException e) {
+      throw new RuntimeException("outputDirURI is not valid - '" + _spec.getOutputDirURI() + "'");
+    }
+    _outputDirFS = PinotFSFactory.create(_outputDirURI.getScheme());
+
+    // Get list of files to process
+    try {
+      _files = _outputDirFS.listFiles(_outputDirURI, true);
+    } catch (IOException e) {
+      throw new RuntimeException("Unable to list all files under outputDirURI - '" + _outputDirURI + "'");
+    }
+  }
+
+  /**
+   * Populates either _segmentsToPush or _segmentUriToTarPathMap based on push type.
+   */
+  public abstract void getSegmentsToPush();
+
+  /**
+   * Returns segmentsTo based on segments obtained from getSegmentsToPush.
+   * The result will to be supplied to the segment replacement protocol when consistent data push is enabled.
+   */
+  public abstract List<String> getSegmentsTo();

Review Comment:
   `getSegmentsToReplace()`



##########
pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/BaseSegmentPushJobRunner.java:
##########
@@ -0,0 +1,142 @@
+/**
+ * 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.plugin.ingestion.batch.common;
+
+import java.io.File;
+import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.segment.local.utils.ConsistentDataPushUtils;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.filesystem.PinotFS;
+import org.apache.pinot.spi.filesystem.PinotFSFactory;
+import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+
+public abstract class BaseSegmentPushJobRunner implements IngestionJobRunner {
+
+  protected SegmentGenerationJobSpec _spec;
+  protected String[] _files;
+  protected PinotFS _outputDirFS;
+  protected URI _outputDirURI;
+  protected List<String> _segmentsToPush = new ArrayList<>();
+  protected Map<String, String> _segmentUriToTarPathMap;
+  protected boolean _consistentPushEnabled;
+
+  /**
+   * Initialize BaseSegmentPushJobRunner with SegmentGenerationJobSpec
+   * Checks for required parameters in the spec and enablement of consistent data push.
+   */
+  @Override
+  public void init(SegmentGenerationJobSpec spec) {
+    _spec = spec;
+    if (_spec.getPushJobSpec() == null) {
+      throw new RuntimeException("Missing PushJobSpec");
+    }
+
+    // Read Table spec
+    if (_spec.getTableSpec() == null) {
+      throw new RuntimeException("Missing tableSpec");
+    }
+
+    // Read Table config
+    if (_spec.getTableSpec().getTableConfigURI() == null) {
+      throw new RuntimeException("Missing property 'tableConfigURI' in 'tableSpec'");
+    }
+
+    _consistentPushEnabled = ConsistentDataPushUtils.consistentDataPushEnabled(_spec);
+  }
+
+  /**
+   * Initialize filesystems and obtain the raw input files for upload.
+   */
+  public void initFileSys() {
+    // init all file systems
+    List<PinotFSSpec> pinotFSSpecs = _spec.getPinotFSSpecs();
+    for (PinotFSSpec pinotFSSpec : pinotFSSpecs) {
+      PinotFSFactory.register(pinotFSSpec.getScheme(), pinotFSSpec.getClassName(), new PinotConfiguration(pinotFSSpec));
+    }
+
+    // Get outputFS for writing output Pinot segments
+    try {
+      _outputDirURI = new URI(_spec.getOutputDirURI());
+      if (_outputDirURI.getScheme() == null) {
+        _outputDirURI = new File(_spec.getOutputDirURI()).toURI();
+      }
+    } catch (URISyntaxException e) {
+      throw new RuntimeException("outputDirURI is not valid - '" + _spec.getOutputDirURI() + "'");
+    }
+    _outputDirFS = PinotFSFactory.create(_outputDirURI.getScheme());
+
+    // Get list of files to process
+    try {
+      _files = _outputDirFS.listFiles(_outputDirURI, true);
+    } catch (IOException e) {
+      throw new RuntimeException("Unable to list all files under outputDirURI - '" + _outputDirURI + "'");
+    }
+  }
+
+  /**
+   * Populates either _segmentsToPush or _segmentUriToTarPathMap based on push type.
+   */
+  public abstract void getSegmentsToPush();
+
+  /**
+   * Returns segmentsTo based on segments obtained from getSegmentsToPush.
+   * The result will to be supplied to the segment replacement protocol when consistent data push is enabled.
+   */
+  public abstract List<String> getSegmentsTo();
+
+  /**
+   * Upload segment obtained by getSegmentsToPush.
+   */
+  public abstract void uploadSegments()
+      throws Exception;
+
+  /**
+   * Runs the main logic of the segment push job runner.
+   * First initialize the filesystem, then upload the segments, while optionally configured to be wrapped around by
+   * the consistent data push protocol.
+   */
+  @Override
+  public void run() {
+    initFileSys();
+    Map<URI, String> uriToLineageEntryIdMap = null;
+    try {
+      getSegmentsToPush();

Review Comment:
   We can change this to
   
   ```
   Map<String, String> segmentsUriToTarPathMap = SegmentPushUtils.getSegmentUriToTarPathMap(_outputDirURI, _pushJobSpec, _files)
   ...
   uploadSegments(segmentsUriToTarPathMap);
   ```
   
   `Map<String, String> segmentsUriToTarPathMap` includes all required information.
   
   1. tarPushJob needs the list of files directly return from `outputDirFS.listFiles()`. This is basically `segmentsUriToTarPathMap.values()`
   
   2. UriPushJob needs the list of Uris and this is the same as `segmentsUriToTarPathMap.keys()`
   
   3. MetadataPush needs the original map.
   
   With this, we don't need to enforce `getSegmentsToPush()` to be a void function with a bit weird requirement to set either `_segmentsToPush` or `_segmentUriToTarPathMap `. We can completely remove the `getSegmentsToPush()` abstract function.
   



##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentUploadIntegrationTest.java:
##########
@@ -188,6 +202,133 @@ public void testUploadAndQuery()
     testCountStar(numDocs);
   }
 
+  /**
+   * Runs both SegmentMetadataPushJobRunner and SegmentTarPushJobRunner while enabling consistent data push.
+   * Checks that segments are properly loaded and segment lineage entry were also in expected states.
+   */
+  @Test
+  public void testUploadAndQueryWithConsistentPush()
+      throws Exception {
+    // Create and upload the schema and table config
+    Schema schema = createSchema();
+    addSchema(schema);
+    TableConfig offlineTableConfig = createOfflineTableConfigWithConsistentPush();
+    addTableConfig(offlineTableConfig);
+
+    List<File> avroFiles = getAllAvroFiles();
+
+    ClusterIntegrationTestUtils.buildSegmentFromAvro(avroFiles.get(0), offlineTableConfig, schema, "_with_move",
+        _segmentDir, _tarDir);
+
+    // First test standalone metadata push job runner
+    BaseSegmentPushJobRunner runner = new SegmentMetadataPushJobRunner();
+    SegmentGenerationJobSpec jobSpec = new SegmentGenerationJobSpec();
+    PushJobSpec pushJobSpec = new PushJobSpec();
+    pushJobSpec.setCopyToDeepStoreForMetadataPush(true);
+    jobSpec.setPushJobSpec(pushJobSpec);
+    PinotFSSpec fsSpec = new PinotFSSpec();
+    fsSpec.setScheme("file");
+    fsSpec.setClassName("org.apache.pinot.spi.filesystem.LocalPinotFS");
+    jobSpec.setPinotFSSpecs(Lists.newArrayList(fsSpec));
+    jobSpec.setOutputDirURI(_tarDir.getAbsolutePath());
+    TableSpec tableSpec = new TableSpec();
+    tableSpec.setTableName(DEFAULT_TABLE_NAME);
+    tableSpec.setTableConfigURI(_controllerRequestURLBuilder.forUpdateTableConfig(DEFAULT_TABLE_NAME));
+    jobSpec.setTableSpec(tableSpec);
+    PinotClusterSpec clusterSpec = new PinotClusterSpec();
+    clusterSpec.setControllerURI(_controllerBaseApiUrl);
+    jobSpec.setPinotClusterSpecs(new PinotClusterSpec[]{clusterSpec});
+
+    File dataDir = new File(_controllerConfig.getDataDir());
+    File dataDirSegments = new File(dataDir, DEFAULT_TABLE_NAME);
+
+    Assert.assertEquals(_tarDir.listFiles().length, 1);
+
+    runner.init(jobSpec);
+    runner.run();
+
+    // Segment should be seen in dataDir
+    Assert.assertTrue(dataDirSegments.exists());
+    Assert.assertEquals(dataDirSegments.listFiles().length, 1);
+    Assert.assertEquals(_tarDir.listFiles().length, 1);
+
+    // test segment loaded
+    JsonNode segmentsList = getSegmentsList();
+    Assert.assertEquals(segmentsList.size(), 1);
+    String segmentNameWithMove = segmentsList.get(0).asText();

Review Comment:
   What do you mean by `segmentNameWithMove` and `segmentNameWithoutMove`? It looks like you are appending the postfix to differentiate segment names. In this case, can we pass `System.currentTimeMillis()` as the postfix and we check whether the segment ends with the timestamp?



##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java:
##########
@@ -0,0 +1,334 @@
+/**
+ * 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.segment.local.utils;
+
+import com.google.common.base.Preconditions;
+import java.io.File;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.common.auth.AuthProviderUtils;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.ingestion.batch.spec.Constants;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class ConsistentDataPushUtils {
+  private ConsistentDataPushUtils() {
+  }
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentPushUtils.class);
+  private static final FileUploadDownloadClient FILE_UPLOAD_DOWNLOAD_CLIENT = new FileUploadDownloadClient();
+  public static final String SEGMENT_NAME_POSTFIX = "segment.name.postfix";
+
+  /**
+   * Checks for enablement of consistent data push. If enabled, fetch the list of segments to be replaced, then
+   * invoke startReplaceSegments API and returns a map of controller URI to lineage entry IDs.
+   * If not, returns an empty hashmap.
+   */
+  public static Map<URI, String> preUpload(SegmentGenerationJobSpec spec, List<String> segmentsTo)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    boolean consistentDataPushEnabled = consistentDataPushEnabled(spec);
+    LOGGER.info("Consistent data push is: {}", consistentDataPushEnabled ? "enabled" : "disabled");
+    Map<URI, String> uriToLineageEntryIdMap = null;
+    if (consistentDataPushEnabled) {
+      LOGGER.info("Start consistent push for table: " + rawTableName);
+      Map<URI, List<String>> uriToExistingOfflineSegments = getSegmentsToReplace(spec, rawTableName);
+      LOGGER.info("Existing segments for table {}: " + uriToExistingOfflineSegments, rawTableName);
+      LOGGER.info("New segments for table: {}: " + segmentsTo, rawTableName);
+      uriToLineageEntryIdMap = startReplaceSegments(spec, uriToExistingOfflineSegments, segmentsTo);
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * uriToLineageEntryIdMap is non-empty if and only if consistent data push is enabled.
+   * If uriToLineageEntryIdMap is non-empty, end the consistent data push protocol for each controller.
+   */
+  public static void postUpload(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap) {
+    String rawTableName = spec.getTableSpec().getTableName();
+    if (uriToLineageEntryIdMap != null && !uriToLineageEntryIdMap.isEmpty()) {
+      LOGGER.info("End consistent push for table: " + rawTableName);
+      endReplaceSegments(spec, uriToLineageEntryIdMap);
+    }
+  }
+
+  /**
+   * Builds a map of controller URI to startReplaceSegments URI for each Pinot cluster in the spec.
+   */
+  public static Map<URI, URI> getStartReplaceSegmentUris(SegmentGenerationJobSpec spec, String rawTableName) {
+    Map<URI, URI> baseUriToStartReplaceSegmentUriMap = new HashMap<>();
+    for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
+      URI controllerURI;
+      try {
+        controllerURI = new URI(pinotClusterSpec.getControllerURI());
+        baseUriToStartReplaceSegmentUriMap.put(controllerURI,
+            FileUploadDownloadClient.getStartReplaceSegmentsURI(controllerURI, rawTableName,
+                TableType.OFFLINE.toString(), true));
+      } catch (URISyntaxException e) {
+        throw new RuntimeException("Got invalid controller uri - '" + pinotClusterSpec.getControllerURI() + "'");
+      }
+    }
+    return baseUriToStartReplaceSegmentUriMap;
+  }
+
+  /**
+   * Starts consistent data push protocol for each Pinot cluster in the spec.
+   * Returns a map of controller URI to segment lineage entry ID.
+   */
+  public static Map<URI, String> startReplaceSegments(SegmentGenerationJobSpec spec,
+      Map<URI, List<String>> uriToSegmentsFrom, List<String> segmentsTo)
+      throws Exception {
+    Map<URI, String> uriToLineageEntryIdMap = new HashMap<>();
+    String rawTableName = spec.getTableSpec().getTableName();
+    Map<URI, URI> segmentsUris = getStartReplaceSegmentUris(spec, rawTableName);
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    LOGGER.info("Start replace segment URIs: " + segmentsUris);
+
+    int attempts = 1;
+    long retryWaitMs = 1000L;
+
+    for (Map.Entry<URI, URI> entry : segmentsUris.entrySet()) {
+      URI controllerUri = entry.getKey();
+      URI startSegmentUri = entry.getValue();
+      List<String> segmentsFrom = uriToSegmentsFrom.get(controllerUri);
+
+      if (!Collections.disjoint(segmentsFrom, segmentsTo)) {
+        String errorMsg =
+            String.format("Found same segment names when attempting to enable consistent push for table: %s",
+                rawTableName);
+        LOGGER.error("SegmentsFrom: {}", segmentsFrom);
+        LOGGER.error("SegmentsTo: {}", segmentsTo);
+        LOGGER.error(errorMsg);
+        throw new RuntimeException(errorMsg);
+      }
+
+      StartReplaceSegmentsRequest startReplaceSegmentsRequest =
+          new StartReplaceSegmentsRequest(segmentsFrom, segmentsTo);
+      RetryPolicies.exponentialBackoffRetryPolicy(attempts, retryWaitMs, 5).attempt(() -> {
+        try {
+          SimpleHttpResponse response =
+              FILE_UPLOAD_DOWNLOAD_CLIENT.startReplaceSegments(startSegmentUri, startReplaceSegmentsRequest,
+                  authProvider);
+
+          String responseString = response.getResponse();
+          LOGGER.info(
+              "Got response {}: {} while sending start replace segment request for table: {}, uploadURI: {}, request:"
+                  + " {}", response.getStatusCode(), responseString, rawTableName, startSegmentUri,
+              startReplaceSegmentsRequest);
+          String segmentLineageEntryId =
+              JsonUtils.stringToJsonNode(responseString).get("segmentLineageEntryId").asText();
+          uriToLineageEntryIdMap.put(controllerUri, segmentLineageEntryId);
+          return true;
+        } catch (SocketTimeoutException se) {
+          // In case of the timeout, we should re-try.
+          return false;
+        } catch (HttpErrorStatusException e) {
+          if (e.getStatusCode() >= 500) {
+            return false;
+          } else {
+            if (e.getStatusCode() == Response.Status.NOT_FOUND.getStatusCode()) {
+              LOGGER.error("Table: {} not found when sending request: {}", rawTableName, startSegmentUri);
+            }
+            throw e;
+          }
+        }
+      });
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * Ends consistent data push protocol for each Pinot cluster in the spec.
+   */
+  public static void endReplaceSegments(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap) {
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    String rawTableName = spec.getTableSpec().getTableName();
+    for (URI uri : uriToLineageEntryIdMap.keySet()) {
+      String segmentLineageEntryId = uriToLineageEntryIdMap.get(uri);
+      try {
+        FILE_UPLOAD_DOWNLOAD_CLIENT.endReplaceSegments(
+            FileUploadDownloadClient.getEndReplaceSegmentsURI(uri, rawTableName, TableType.OFFLINE.toString(),
+                segmentLineageEntryId), HttpClient.DEFAULT_SOCKET_TIMEOUT_MS, authProvider);
+      } catch (URISyntaxException e) {
+        throw new RuntimeException("Got invalid controller uri - '" + uri + "'");
+      } catch (HttpErrorStatusException | IOException e) {
+        e.printStackTrace();
+      }
+    }
+  }
+
+  /**
+   * Revert segment lineage entry when exception gets caught. This revert request is called at best effort.
+   * If the revert call fails at this point, the next startReplaceSegment call will do the cleanup
+   * by marking the previous entry to "REVERTED" and cleaning up the leftover segments.
+   */
+  public static void handleUploadException(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap,
+      Exception exception) {
+    if (uriToLineageEntryIdMap != null) {
+      LOGGER.error("Exception when pushing segments. Marking segment lineage entry to 'REVERTED'.", exception);
+      String rawTableName = spec.getTableSpec().getTableName();
+      for (Map.Entry<URI, String> entry : uriToLineageEntryIdMap.entrySet()) {
+        String segmentLineageEntryId = entry.getValue();
+        try {
+          URI uri = FileUploadDownloadClient.getRevertReplaceSegmentsURI(entry.getKey(), rawTableName,
+              TableType.OFFLINE.name(), segmentLineageEntryId, true);
+          SimpleHttpResponse response = FILE_UPLOAD_DOWNLOAD_CLIENT.revertReplaceSegments(uri);
+          LOGGER.info("Got response {}: {} while sending revert replace segment request for table: {}, uploadURI: {}",
+              response.getStatusCode(), response.getResponse(), rawTableName, entry.getKey());
+        } catch (URISyntaxException | HttpErrorStatusException | IOException e) {
+          LOGGER.error("Exception when sending revert replace segment request to controller: {} for table: {}",
+              entry.getKey(), rawTableName, e);
+        }
+      }
+    }
+  }
+
+  /**
+   * Ensures that all files in tarFilePaths have the expected tar file extension and obtain segment names given
+   * tarFilePaths.
+   */
+  public static List<String> getTarSegmentsTo(List<String> tarFilePaths) {
+    List<String> segmentsTo = new ArrayList<>();
+    for (String tarFilePath : tarFilePaths) {
+      File tarFile = new File(tarFilePath);
+      String fileName = tarFile.getName();
+      Preconditions.checkArgument(fileName.endsWith(Constants.TAR_GZ_FILE_EXT));
+      String segmentName = getSegmentNameFromFilePath(fileName);
+      segmentsTo.add(segmentName);
+    }
+    return segmentsTo;
+  }
+
+  /**
+   * Ensures that all URIs in segmentUris have the expected tar file extension and obtain segment names given
+   * segmentUris.
+   */
+  public static List<String> getUriSegmentsTo(List<String> segmentUris) {
+    List<String> segmentsTo = new ArrayList<>();
+    for (String segmentUri : segmentUris) {
+      Preconditions.checkArgument(segmentUri.endsWith(Constants.TAR_GZ_FILE_EXT));
+      String segmentName = getSegmentNameFromFilePath(segmentUri);
+      segmentsTo.add(segmentName);
+    }
+    return segmentsTo;
+  }
+
+  /**
+   * Ensures that all tarPaths in segmentUriToTarPathMap have the expected tar file extension and obtain segment names
+   * given tarPaths.
+   */
+  public static List<String> getMetadataSegmentsTo(Map<String, String> segmentUriToTarPathMap) {

Review Comment:
   What does it mean by `Metadata` 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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] jtao15 commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
jtao15 commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r960232443


##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java:
##########
@@ -0,0 +1,296 @@
+/**
+ * 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.segment.local.utils;
+
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.common.auth.AuthProviderUtils;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.apache.pinot.spi.utils.retry.RetryPolicy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class ConsistentDataPushUtils {
+  private ConsistentDataPushUtils() {
+  }
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentPushUtils.class);
+  private static final FileUploadDownloadClient FILE_UPLOAD_DOWNLOAD_CLIENT = new FileUploadDownloadClient();
+  private static final RetryPolicy DEFAULT_RETRY_POLICY = RetryPolicies.exponentialBackoffRetryPolicy(5, 10_000L, 2.0);
+  public static final String SEGMENT_NAME_POSTFIX = "segment.name.postfix";
+
+  /**
+   * Checks for enablement of consistent data push. If enabled, fetch the list of segments to be replaced, then
+   * invoke startReplaceSegments API and returns a map of controller URI to lineage entry IDs.
+   * If not, returns an empty hashmap.
+   */
+  public static Map<URI, String> preUpload(SegmentGenerationJobSpec spec, List<String> segmentsTo)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    boolean consistentDataPushEnabled = consistentDataPushEnabled(spec);
+    LOGGER.info("Consistent data push is: {}", consistentDataPushEnabled ? "enabled" : "disabled");
+    Map<URI, String> uriToLineageEntryIdMap = null;
+    if (consistentDataPushEnabled) {
+      LOGGER.info("Start consistent push for table: " + rawTableName);
+      Map<URI, List<String>> uriToExistingOfflineSegments = getSegmentsToReplace(spec, rawTableName);
+      LOGGER.info("Existing segments for table {}: " + uriToExistingOfflineSegments, rawTableName);
+      LOGGER.info("New segments for table: {}: " + segmentsTo, rawTableName);
+      uriToLineageEntryIdMap = startReplaceSegments(spec, uriToExistingOfflineSegments, segmentsTo);
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * uriToLineageEntryIdMap is non-empty if and only if consistent data push is enabled.
+   * If uriToLineageEntryIdMap is non-empty, end the consistent data push protocol for each controller.
+   */
+  public static void postUpload(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    if (uriToLineageEntryIdMap != null && !uriToLineageEntryIdMap.isEmpty()) {
+      LOGGER.info("End consistent push for table: " + rawTableName);
+      endReplaceSegments(spec, uriToLineageEntryIdMap);
+    }
+  }
+
+  /**
+   * Builds a map of controller URI to startReplaceSegments URI for each Pinot cluster in the spec.
+   */
+  public static Map<URI, URI> getStartReplaceSegmentUris(SegmentGenerationJobSpec spec, String rawTableName) {
+    Map<URI, URI> baseUriToStartReplaceSegmentUriMap = new HashMap<>();
+    for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
+      URI controllerURI;
+      try {
+        controllerURI = new URI(pinotClusterSpec.getControllerURI());
+        baseUriToStartReplaceSegmentUriMap.put(controllerURI,
+            FileUploadDownloadClient.getStartReplaceSegmentsURI(controllerURI, rawTableName,
+                TableType.OFFLINE.toString(), true));
+      } catch (URISyntaxException e) {
+        throw new RuntimeException("Got invalid controller uri - '" + pinotClusterSpec.getControllerURI() + "'");
+      }
+    }
+    return baseUriToStartReplaceSegmentUriMap;
+  }
+
+  /**
+   * Starts consistent data push protocol for each Pinot cluster in the spec.
+   * Returns a map of controller URI to segment lineage entry ID.
+   */
+  public static Map<URI, String> startReplaceSegments(SegmentGenerationJobSpec spec,
+      Map<URI, List<String>> uriToSegmentsFrom, List<String> segmentsTo)
+      throws Exception {
+    Map<URI, String> uriToLineageEntryIdMap = new HashMap<>();
+    String rawTableName = spec.getTableSpec().getTableName();
+    Map<URI, URI> segmentsUris = getStartReplaceSegmentUris(spec, rawTableName);
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    LOGGER.info("Start replace segment URIs: " + segmentsUris);
+
+    for (Map.Entry<URI, URI> entry : segmentsUris.entrySet()) {
+      URI controllerUri = entry.getKey();
+      URI startSegmentUri = entry.getValue();
+      List<String> segmentsFrom = uriToSegmentsFrom.get(controllerUri);
+
+      if (!Collections.disjoint(segmentsFrom, segmentsTo)) {

Review Comment:
   We are checking the table contains all segments in `segmentsFrom`, and doesn't have any segment in `segmentsTo` in startReplaceSegments(). This check is not needed?



##########
pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java:
##########
@@ -747,6 +754,72 @@ public SimpleHttpResponse uploadSegment(URI uri, String segmentName, InputStream
     return uploadSegment(uri, segmentName, inputStream, null, parameters, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
   }
 
+  /**
+   * Returns a map from a given tableType to a list of segments for that given tableType (OFFLINE or REALTIME)
+   * If tableType is left unspecified, both OFFLINE and REALTIME segments will be returned in the map.
+   */
+  public Map<String, List<String>> getSegments(URI controllerUri, String rawTableName, @Nullable TableType tableType,
+      boolean excludeReplacedSegments) throws Exception {
+    List<String> tableTypes;
+    if (tableType == null) {
+      tableTypes = Arrays.asList(TableType.OFFLINE.toString(), TableType.REALTIME.toString());
+    } else {
+      tableTypes = Arrays.asList(tableType.toString());
+    }
+    ControllerRequestURLBuilder controllerRequestURLBuilder =
+        ControllerRequestURLBuilder.baseUrl(controllerUri.toString());
+    Map<String, List<String>> tableTypeToSegments = new HashMap<>();
+    for (String tableTypeToFilter : tableTypes) {
+      tableTypeToSegments.put(tableTypeToFilter, new ArrayList<>());
+      String uri = controllerRequestURLBuilder.forSegmentListAPIWithTableTypeAndExcludeReplacedSegments(rawTableName,
+          tableTypeToFilter, excludeReplacedSegments);
+      RequestBuilder requestBuilder = RequestBuilder.get(uri).setVersion(HttpVersion.HTTP_1_1);
+      HttpClient.setTimeout(requestBuilder, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
+      RetryPolicies.exponentialBackoffRetryPolicy(5, 10_000L, 2.0).attempt(() -> {
+        try {
+          SimpleHttpResponse response =
+              HttpClient.wrapAndThrowHttpException(_httpClient.sendRequest(requestBuilder.build()));
+          LOGGER.info("Response {}: {} received for GET request to URI: {}", response.getStatusCode(),
+              response.getResponse(), uri);
+          tableTypeToSegments.put(tableTypeToFilter,
+              getSegmentNamesFromResponse(tableTypeToFilter, response.getResponse()));
+          return true;
+        } catch (SocketTimeoutException se) {
+          // In case of the timeout, we should re-try.
+          return false;
+        } catch (HttpErrorStatusException e) {
+          if (e.getStatusCode() < 500) {
+            if (e.getStatusCode() == Response.Status.NOT_FOUND.getStatusCode()) {
+              LOGGER.error("Segments not found for table {} when sending request uri: {}", rawTableName, uri);
+            }
+          }
+          return false;
+        }
+      });
+    }
+    return tableTypeToSegments;
+  }
+
+  public List<String> getSegmentNamesFromResponse(String tableType, String responseString)

Review Comment:
   (nit) private?



##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java:
##########
@@ -0,0 +1,296 @@
+/**
+ * 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.segment.local.utils;
+
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.common.auth.AuthProviderUtils;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.apache.pinot.spi.utils.retry.RetryPolicy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class ConsistentDataPushUtils {
+  private ConsistentDataPushUtils() {
+  }
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentPushUtils.class);
+  private static final FileUploadDownloadClient FILE_UPLOAD_DOWNLOAD_CLIENT = new FileUploadDownloadClient();
+  private static final RetryPolicy DEFAULT_RETRY_POLICY = RetryPolicies.exponentialBackoffRetryPolicy(5, 10_000L, 2.0);

Review Comment:
   lgtm



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r959249986


##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java:
##########
@@ -0,0 +1,284 @@
+/**
+ * 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.segment.local.utils;
+
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.common.auth.AuthProviderUtils;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class ConsistentDataPushUtils {
+  private ConsistentDataPushUtils() {
+  }
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentPushUtils.class);
+  private static final FileUploadDownloadClient FILE_UPLOAD_DOWNLOAD_CLIENT = new FileUploadDownloadClient();
+  public static final String SEGMENT_NAME_POSTFIX = "segment.name.postfix";
+
+  /**
+   * Checks for enablement of consistent data push. If enabled, fetch the list of segments to be replaced, then
+   * invoke startReplaceSegments API and returns a map of controller URI to lineage entry IDs.
+   * If not, returns an empty hashmap.
+   */
+  public static Map<URI, String> preUpload(SegmentGenerationJobSpec spec, List<String> segmentsTo)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    boolean consistentDataPushEnabled = consistentDataPushEnabled(spec);
+    LOGGER.info("Consistent data push is: {}", consistentDataPushEnabled ? "enabled" : "disabled");
+    Map<URI, String> uriToLineageEntryIdMap = null;
+    if (consistentDataPushEnabled) {
+      LOGGER.info("Start consistent push for table: " + rawTableName);
+      Map<URI, List<String>> uriToExistingOfflineSegments = getSegmentsToReplace(spec, rawTableName);
+      LOGGER.info("Existing segments for table {}: " + uriToExistingOfflineSegments, rawTableName);
+      LOGGER.info("New segments for table: {}: " + segmentsTo, rawTableName);
+      uriToLineageEntryIdMap = startReplaceSegments(spec, uriToExistingOfflineSegments, segmentsTo);
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * uriToLineageEntryIdMap is non-empty if and only if consistent data push is enabled.
+   * If uriToLineageEntryIdMap is non-empty, end the consistent data push protocol for each controller.
+   */
+  public static void postUpload(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap) {
+    String rawTableName = spec.getTableSpec().getTableName();
+    if (uriToLineageEntryIdMap != null && !uriToLineageEntryIdMap.isEmpty()) {
+      LOGGER.info("End consistent push for table: " + rawTableName);
+      endReplaceSegments(spec, uriToLineageEntryIdMap);
+    }
+  }
+
+  /**
+   * Builds a map of controller URI to startReplaceSegments URI for each Pinot cluster in the spec.
+   */
+  public static Map<URI, URI> getStartReplaceSegmentUris(SegmentGenerationJobSpec spec, String rawTableName) {
+    Map<URI, URI> baseUriToStartReplaceSegmentUriMap = new HashMap<>();
+    for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
+      URI controllerURI;
+      try {
+        controllerURI = new URI(pinotClusterSpec.getControllerURI());
+        baseUriToStartReplaceSegmentUriMap.put(controllerURI,
+            FileUploadDownloadClient.getStartReplaceSegmentsURI(controllerURI, rawTableName,
+                TableType.OFFLINE.toString(), true));
+      } catch (URISyntaxException e) {
+        throw new RuntimeException("Got invalid controller uri - '" + pinotClusterSpec.getControllerURI() + "'");
+      }
+    }
+    return baseUriToStartReplaceSegmentUriMap;
+  }
+
+  /**
+   * Starts consistent data push protocol for each Pinot cluster in the spec.
+   * Returns a map of controller URI to segment lineage entry ID.
+   */
+  public static Map<URI, String> startReplaceSegments(SegmentGenerationJobSpec spec,
+      Map<URI, List<String>> uriToSegmentsFrom, List<String> segmentsTo)
+      throws Exception {
+    Map<URI, String> uriToLineageEntryIdMap = new HashMap<>();
+    String rawTableName = spec.getTableSpec().getTableName();
+    Map<URI, URI> segmentsUris = getStartReplaceSegmentUris(spec, rawTableName);
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    LOGGER.info("Start replace segment URIs: " + segmentsUris);
+
+    int attempts = 1;
+    long retryWaitMs = 1000L;
+
+    for (Map.Entry<URI, URI> entry : segmentsUris.entrySet()) {
+      URI controllerUri = entry.getKey();
+      URI startSegmentUri = entry.getValue();
+      List<String> segmentsFrom = uriToSegmentsFrom.get(controllerUri);
+
+      if (!Collections.disjoint(segmentsFrom, segmentsTo)) {
+        String errorMsg =
+            String.format("Found same segment names when attempting to enable consistent push for table: %s",
+                rawTableName);
+        LOGGER.error("SegmentsFrom: {}", segmentsFrom);
+        LOGGER.error("SegmentsTo: {}", segmentsTo);
+        LOGGER.error(errorMsg);
+        throw new RuntimeException(errorMsg);
+      }
+
+      StartReplaceSegmentsRequest startReplaceSegmentsRequest =
+          new StartReplaceSegmentsRequest(segmentsFrom, segmentsTo);
+      RetryPolicies.exponentialBackoffRetryPolicy(attempts, retryWaitMs, 5).attempt(() -> {
+        try {
+          SimpleHttpResponse response =
+              FILE_UPLOAD_DOWNLOAD_CLIENT.startReplaceSegments(startSegmentUri, startReplaceSegmentsRequest,
+                  authProvider);
+
+          String responseString = response.getResponse();
+          LOGGER.info(
+              "Got response {}: {} while sending start replace segment request for table: {}, uploadURI: {}, request:"
+                  + " {}", response.getStatusCode(), responseString, rawTableName, startSegmentUri,
+              startReplaceSegmentsRequest);
+          String segmentLineageEntryId =
+              JsonUtils.stringToJsonNode(responseString).get("segmentLineageEntryId").asText();
+          uriToLineageEntryIdMap.put(controllerUri, segmentLineageEntryId);
+          return true;
+        } catch (SocketTimeoutException se) {
+          // In case of the timeout, we should re-try.
+          return false;
+        } catch (HttpErrorStatusException e) {
+          if (e.getStatusCode() >= 500) {
+            return false;
+          } else {
+            if (e.getStatusCode() == Response.Status.NOT_FOUND.getStatusCode()) {
+              LOGGER.error("Table: {} not found when sending request: {}", rawTableName, startSegmentUri);
+            }
+            throw e;
+          }
+        }
+      });
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * Ends consistent data push protocol for each Pinot cluster in the spec.
+   */
+  public static void endReplaceSegments(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap) {
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    String rawTableName = spec.getTableSpec().getTableName();
+    for (URI uri : uriToLineageEntryIdMap.keySet()) {
+      String segmentLineageEntryId = uriToLineageEntryIdMap.get(uri);
+      try {
+        FILE_UPLOAD_DOWNLOAD_CLIENT.endReplaceSegments(

Review Comment:
   can we add the retry here as well to keep the consistency with the startReplaceSegments call?



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r958990524


##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java:
##########
@@ -0,0 +1,334 @@
+/**
+ * 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.segment.local.utils;
+
+import com.google.common.base.Preconditions;
+import java.io.File;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.common.auth.AuthProviderUtils;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.ingestion.batch.spec.Constants;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class ConsistentDataPushUtils {
+  private ConsistentDataPushUtils() {
+  }
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentPushUtils.class);
+  private static final FileUploadDownloadClient FILE_UPLOAD_DOWNLOAD_CLIENT = new FileUploadDownloadClient();
+  public static final String SEGMENT_NAME_POSTFIX = "segment.name.postfix";
+
+  /**
+   * Checks for enablement of consistent data push. If enabled, fetch the list of segments to be replaced, then
+   * invoke startReplaceSegments API and returns a map of controller URI to lineage entry IDs.
+   * If not, returns an empty hashmap.
+   */
+  public static Map<URI, String> preUpload(SegmentGenerationJobSpec spec, List<String> segmentsTo)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    boolean consistentDataPushEnabled = consistentDataPushEnabled(spec);
+    LOGGER.info("Consistent data push is: {}", consistentDataPushEnabled ? "enabled" : "disabled");
+    Map<URI, String> uriToLineageEntryIdMap = null;
+    if (consistentDataPushEnabled) {
+      LOGGER.info("Start consistent push for table: " + rawTableName);
+      Map<URI, List<String>> uriToExistingOfflineSegments = getSegmentsToReplace(spec, rawTableName);
+      LOGGER.info("Existing segments for table {}: " + uriToExistingOfflineSegments, rawTableName);
+      LOGGER.info("New segments for table: {}: " + segmentsTo, rawTableName);
+      uriToLineageEntryIdMap = startReplaceSegments(spec, uriToExistingOfflineSegments, segmentsTo);
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * uriToLineageEntryIdMap is non-empty if and only if consistent data push is enabled.
+   * If uriToLineageEntryIdMap is non-empty, end the consistent data push protocol for each controller.
+   */
+  public static void postUpload(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap) {
+    String rawTableName = spec.getTableSpec().getTableName();
+    if (uriToLineageEntryIdMap != null && !uriToLineageEntryIdMap.isEmpty()) {
+      LOGGER.info("End consistent push for table: " + rawTableName);
+      endReplaceSegments(spec, uriToLineageEntryIdMap);
+    }
+  }
+
+  /**
+   * Builds a map of controller URI to startReplaceSegments URI for each Pinot cluster in the spec.
+   */
+  public static Map<URI, URI> getStartReplaceSegmentUris(SegmentGenerationJobSpec spec, String rawTableName) {
+    Map<URI, URI> baseUriToStartReplaceSegmentUriMap = new HashMap<>();
+    for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
+      URI controllerURI;
+      try {
+        controllerURI = new URI(pinotClusterSpec.getControllerURI());
+        baseUriToStartReplaceSegmentUriMap.put(controllerURI,
+            FileUploadDownloadClient.getStartReplaceSegmentsURI(controllerURI, rawTableName,
+                TableType.OFFLINE.toString(), true));
+      } catch (URISyntaxException e) {
+        throw new RuntimeException("Got invalid controller uri - '" + pinotClusterSpec.getControllerURI() + "'");
+      }
+    }
+    return baseUriToStartReplaceSegmentUriMap;
+  }
+
+  /**
+   * Starts consistent data push protocol for each Pinot cluster in the spec.
+   * Returns a map of controller URI to segment lineage entry ID.
+   */
+  public static Map<URI, String> startReplaceSegments(SegmentGenerationJobSpec spec,
+      Map<URI, List<String>> uriToSegmentsFrom, List<String> segmentsTo)
+      throws Exception {
+    Map<URI, String> uriToLineageEntryIdMap = new HashMap<>();
+    String rawTableName = spec.getTableSpec().getTableName();
+    Map<URI, URI> segmentsUris = getStartReplaceSegmentUris(spec, rawTableName);
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    LOGGER.info("Start replace segment URIs: " + segmentsUris);
+
+    int attempts = 1;
+    long retryWaitMs = 1000L;
+
+    for (Map.Entry<URI, URI> entry : segmentsUris.entrySet()) {
+      URI controllerUri = entry.getKey();
+      URI startSegmentUri = entry.getValue();
+      List<String> segmentsFrom = uriToSegmentsFrom.get(controllerUri);
+
+      if (!Collections.disjoint(segmentsFrom, segmentsTo)) {
+        String errorMsg =
+            String.format("Found same segment names when attempting to enable consistent push for table: %s",
+                rawTableName);
+        LOGGER.error("SegmentsFrom: {}", segmentsFrom);
+        LOGGER.error("SegmentsTo: {}", segmentsTo);
+        LOGGER.error(errorMsg);
+        throw new RuntimeException(errorMsg);
+      }
+
+      StartReplaceSegmentsRequest startReplaceSegmentsRequest =
+          new StartReplaceSegmentsRequest(segmentsFrom, segmentsTo);
+      RetryPolicies.exponentialBackoffRetryPolicy(attempts, retryWaitMs, 5).attempt(() -> {
+        try {
+          SimpleHttpResponse response =
+              FILE_UPLOAD_DOWNLOAD_CLIENT.startReplaceSegments(startSegmentUri, startReplaceSegmentsRequest,
+                  authProvider);
+
+          String responseString = response.getResponse();
+          LOGGER.info(
+              "Got response {}: {} while sending start replace segment request for table: {}, uploadURI: {}, request:"
+                  + " {}", response.getStatusCode(), responseString, rawTableName, startSegmentUri,
+              startReplaceSegmentsRequest);
+          String segmentLineageEntryId =
+              JsonUtils.stringToJsonNode(responseString).get("segmentLineageEntryId").asText();
+          uriToLineageEntryIdMap.put(controllerUri, segmentLineageEntryId);
+          return true;
+        } catch (SocketTimeoutException se) {
+          // In case of the timeout, we should re-try.
+          return false;
+        } catch (HttpErrorStatusException e) {
+          if (e.getStatusCode() >= 500) {
+            return false;
+          } else {
+            if (e.getStatusCode() == Response.Status.NOT_FOUND.getStatusCode()) {
+              LOGGER.error("Table: {} not found when sending request: {}", rawTableName, startSegmentUri);
+            }
+            throw e;
+          }
+        }
+      });
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * Ends consistent data push protocol for each Pinot cluster in the spec.
+   */
+  public static void endReplaceSegments(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap) {
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    String rawTableName = spec.getTableSpec().getTableName();
+    for (URI uri : uriToLineageEntryIdMap.keySet()) {
+      String segmentLineageEntryId = uriToLineageEntryIdMap.get(uri);
+      try {
+        FILE_UPLOAD_DOWNLOAD_CLIENT.endReplaceSegments(
+            FileUploadDownloadClient.getEndReplaceSegmentsURI(uri, rawTableName, TableType.OFFLINE.toString(),
+                segmentLineageEntryId), HttpClient.DEFAULT_SOCKET_TIMEOUT_MS, authProvider);
+      } catch (URISyntaxException e) {
+        throw new RuntimeException("Got invalid controller uri - '" + uri + "'");
+      } catch (HttpErrorStatusException | IOException e) {
+        e.printStackTrace();
+      }
+    }
+  }
+
+  /**
+   * Revert segment lineage entry when exception gets caught. This revert request is called at best effort.
+   * If the revert call fails at this point, the next startReplaceSegment call will do the cleanup
+   * by marking the previous entry to "REVERTED" and cleaning up the leftover segments.
+   */
+  public static void handleUploadException(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap,
+      Exception exception) {
+    if (uriToLineageEntryIdMap != null) {
+      LOGGER.error("Exception when pushing segments. Marking segment lineage entry to 'REVERTED'.", exception);
+      String rawTableName = spec.getTableSpec().getTableName();
+      for (Map.Entry<URI, String> entry : uriToLineageEntryIdMap.entrySet()) {
+        String segmentLineageEntryId = entry.getValue();
+        try {
+          URI uri = FileUploadDownloadClient.getRevertReplaceSegmentsURI(entry.getKey(), rawTableName,
+              TableType.OFFLINE.name(), segmentLineageEntryId, true);
+          SimpleHttpResponse response = FILE_UPLOAD_DOWNLOAD_CLIENT.revertReplaceSegments(uri);
+          LOGGER.info("Got response {}: {} while sending revert replace segment request for table: {}, uploadURI: {}",
+              response.getStatusCode(), response.getResponse(), rawTableName, entry.getKey());
+        } catch (URISyntaxException | HttpErrorStatusException | IOException e) {
+          LOGGER.error("Exception when sending revert replace segment request to controller: {} for table: {}",
+              entry.getKey(), rawTableName, e);
+        }
+      }
+    }
+  }
+
+  /**
+   * Ensures that all files in tarFilePaths have the expected tar file extension and obtain segment names given
+   * tarFilePaths.
+   */
+  public static List<String> getTarSegmentsTo(List<String> tarFilePaths) {
+    List<String> segmentsTo = new ArrayList<>();
+    for (String tarFilePath : tarFilePaths) {
+      File tarFile = new File(tarFilePath);
+      String fileName = tarFile.getName();
+      Preconditions.checkArgument(fileName.endsWith(Constants.TAR_GZ_FILE_EXT));
+      String segmentName = getSegmentNameFromFilePath(fileName);
+      segmentsTo.add(segmentName);
+    }
+    return segmentsTo;
+  }
+
+  /**
+   * Ensures that all URIs in segmentUris have the expected tar file extension and obtain segment names given
+   * segmentUris.
+   */
+  public static List<String> getUriSegmentsTo(List<String> segmentUris) {
+    List<String> segmentsTo = new ArrayList<>();
+    for (String segmentUri : segmentUris) {
+      Preconditions.checkArgument(segmentUri.endsWith(Constants.TAR_GZ_FILE_EXT));
+      String segmentName = getSegmentNameFromFilePath(segmentUri);
+      segmentsTo.add(segmentName);
+    }
+    return segmentsTo;
+  }
+
+  /**
+   * Ensures that all tarPaths in segmentUriToTarPathMap have the expected tar file extension and obtain segment names
+   * given tarPaths.
+   */
+  public static List<String> getMetadataSegmentsTo(Map<String, String> segmentUriToTarPathMap) {

Review Comment:
   `getSegmentsToFromMetadata` would've been a better name, but refactored this area to make code more readable.



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r959049111


##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java:
##########
@@ -371,4 +372,36 @@ private static File generateSegmentMetadataFile(PinotFS fileSystem, URI tarFileU
       FileUtils.deleteQuietly(segmentMetadataDir);
     }
   }
+
+  public static List<String> getSegmentNames(BatchConfigProperties.SegmentPushType pushMode,
+      Map<String, String> segmentsUriToTarPathMap) {
+    List<String> segmentNames = new ArrayList<>();
+    if (pushMode.equals(BatchConfigProperties.SegmentPushType.TAR) || pushMode.equals(
+        BatchConfigProperties.SegmentPushType.METADATA)) {
+      for (String tarFilePath : segmentsUriToTarPathMap.values()) {
+        File tarFile = new File(tarFilePath);
+        String fileName = tarFile.getName();
+        Preconditions.checkArgument(fileName.endsWith(Constants.TAR_GZ_FILE_EXT));
+        String segmentName = getSegmentNameFromPath(fileName);
+        segmentNames.add(segmentName);
+      }
+    }
+    if (pushMode.equals(BatchConfigProperties.SegmentPushType.URI)) {
+      for (String segmentUri : segmentsUriToTarPathMap.keySet()) {
+        Preconditions.checkArgument(segmentUri.endsWith(Constants.TAR_GZ_FILE_EXT));
+        String segmentName = getSegmentNameFromPath(segmentUri);
+        segmentNames.add(segmentName);
+      }
+    }
+    return segmentNames;
+  }
+
+  /**
+   * Obtain segment name given filePath by reading from after the last slash (if present) up to and before the tar
+   * extension.
+   */
+  public static String getSegmentNameFromPath(String filePath) {
+    int startIndex = filePath.contains("/") ? filePath.lastIndexOf("/") + 1 : 0;

Review Comment:
   Thanks for the reminder!



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r964300002


##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java:
##########
@@ -371,4 +372,35 @@ private static File generateSegmentMetadataFile(PinotFS fileSystem, URI tarFileU
       FileUtils.deleteQuietly(segmentMetadataDir);
     }
   }
+
+  public static List<String> getSegmentNames(BatchConfigProperties.SegmentPushType pushMode,

Review Comment:
   Thanks @snleee! As discussed offline, after looking at `SegmentPushUtils.java#getSegmentUriToTarPathMap` again, this separation based on `pushMode` is indeed unnecessary as the URI names are derived from tarFilePaths from `segmentsUriToTarPathMap.values()` anyways, which means TAR, METADATA and URI can be unified.
   
   Hence, removed the overrides and brought generic implementation up to the base class.



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r964338431


##########
pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-standalone/src/main/java/org/apache/pinot/plugin/ingestion/batch/standalone/SegmentTarPushJobRunner.java:
##########
@@ -46,48 +35,8 @@ public SegmentTarPushJobRunner(SegmentGenerationJobSpec spec) {
     init(spec);
   }
 
-  @Override
-  public void init(SegmentGenerationJobSpec spec) {
-    _spec = spec;
-  }
-
-  @Override
-  public void run() {
-    //init all file systems
-    List<PinotFSSpec> pinotFSSpecs = _spec.getPinotFSSpecs();
-    for (PinotFSSpec pinotFSSpec : pinotFSSpecs) {
-      PinotFSFactory.register(pinotFSSpec.getScheme(), pinotFSSpec.getClassName(), new PinotConfiguration(pinotFSSpec));
-    }
-
-    //Get outputFS for writing output pinot segments
-    URI outputDirURI;
-    try {
-      outputDirURI = new URI(_spec.getOutputDirURI());
-      if (outputDirURI.getScheme() == null) {
-        outputDirURI = new File(_spec.getOutputDirURI()).toURI();
-      }
-    } catch (URISyntaxException e) {
-      throw new RuntimeException("outputDirURI is not valid - '" + _spec.getOutputDirURI() + "'");
-    }
-    PinotFS outputDirFS = PinotFSFactory.create(outputDirURI.getScheme());
-    //Get list of files to process
-    String[] files;
-    try {
-      files = outputDirFS.listFiles(outputDirURI, true);
-    } catch (IOException e) {
-      throw new RuntimeException("Unable to list all files under outputDirURI - '" + outputDirURI + "'");
-    }
-
-    List<String> segmentsToPush = new ArrayList<>();
-    for (String file : files) {
-      if (file.endsWith(Constants.TAR_GZ_FILE_EXT)) {
-        segmentsToPush.add(file);
-      }
-    }
-    try {
-      SegmentPushUtils.pushSegments(_spec, outputDirFS, segmentsToPush);
-    } catch (RetriableOperationException | AttemptsExceededException e) {
-      throw new RuntimeException(e);
-    }
+  public void uploadSegments(Map<String, String> segmentsUriToTarPathMap)
+      throws AttemptsExceededException, RetriableOperationException {
+    SegmentPushUtils.pushSegments(_spec, _outputDirFS, new ArrayList<>(segmentsUriToTarPathMap.values()));

Review Comment:
   For copying the elements from `map.values()` to `List`? 
   
   In our case, the worst-case scenario should be 10-100 thousand and I think it's OK to keep it this way. I don't think that it's a performance-critical path.



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r964338431


##########
pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-standalone/src/main/java/org/apache/pinot/plugin/ingestion/batch/standalone/SegmentTarPushJobRunner.java:
##########
@@ -46,48 +35,8 @@ public SegmentTarPushJobRunner(SegmentGenerationJobSpec spec) {
     init(spec);
   }
 
-  @Override
-  public void init(SegmentGenerationJobSpec spec) {
-    _spec = spec;
-  }
-
-  @Override
-  public void run() {
-    //init all file systems
-    List<PinotFSSpec> pinotFSSpecs = _spec.getPinotFSSpecs();
-    for (PinotFSSpec pinotFSSpec : pinotFSSpecs) {
-      PinotFSFactory.register(pinotFSSpec.getScheme(), pinotFSSpec.getClassName(), new PinotConfiguration(pinotFSSpec));
-    }
-
-    //Get outputFS for writing output pinot segments
-    URI outputDirURI;
-    try {
-      outputDirURI = new URI(_spec.getOutputDirURI());
-      if (outputDirURI.getScheme() == null) {
-        outputDirURI = new File(_spec.getOutputDirURI()).toURI();
-      }
-    } catch (URISyntaxException e) {
-      throw new RuntimeException("outputDirURI is not valid - '" + _spec.getOutputDirURI() + "'");
-    }
-    PinotFS outputDirFS = PinotFSFactory.create(outputDirURI.getScheme());
-    //Get list of files to process
-    String[] files;
-    try {
-      files = outputDirFS.listFiles(outputDirURI, true);
-    } catch (IOException e) {
-      throw new RuntimeException("Unable to list all files under outputDirURI - '" + outputDirURI + "'");
-    }
-
-    List<String> segmentsToPush = new ArrayList<>();
-    for (String file : files) {
-      if (file.endsWith(Constants.TAR_GZ_FILE_EXT)) {
-        segmentsToPush.add(file);
-      }
-    }
-    try {
-      SegmentPushUtils.pushSegments(_spec, outputDirFS, segmentsToPush);
-    } catch (RetriableOperationException | AttemptsExceededException e) {
-      throw new RuntimeException(e);
-    }
+  public void uploadSegments(Map<String, String> segmentsUriToTarPathMap)
+      throws AttemptsExceededException, RetriableOperationException {
+    SegmentPushUtils.pushSegments(_spec, _outputDirFS, new ArrayList<>(segmentsUriToTarPathMap.values()));

Review Comment:
   For copying the elements from `map.values()` to `List`? 
   
   In our case, the worst-case scenario should be 10s thousands and I think it's OK to keep it this way. I don't think that it's a performance-critical path.



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r964301549


##########
pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-standalone/src/main/java/org/apache/pinot/plugin/ingestion/batch/standalone/SegmentUriPushJobRunner.java:
##########
@@ -47,55 +36,8 @@ public SegmentUriPushJobRunner(SegmentGenerationJobSpec spec) {
     init(spec);
   }
 
-  @Override
-  public void init(SegmentGenerationJobSpec spec) {
-    _spec = spec;
-    if (_spec.getPushJobSpec() == null) {
-      throw new RuntimeException("Missing PushJobSpec");
-    }
-  }
-
-  @Override
-  public void run() {
-    //init all file systems
-    List<PinotFSSpec> pinotFSSpecs = _spec.getPinotFSSpecs();
-    for (PinotFSSpec pinotFSSpec : pinotFSSpecs) {
-      PinotFSFactory.register(pinotFSSpec.getScheme(), pinotFSSpec.getClassName(), new PinotConfiguration(pinotFSSpec));
-    }
-
-    //Get outputFS for writing output Pinot segments
-    URI outputDirURI;
-    try {
-      outputDirURI = new URI(_spec.getOutputDirURI());
-      if (outputDirURI.getScheme() == null) {
-        outputDirURI = new File(_spec.getOutputDirURI()).toURI();
-      }
-    } catch (URISyntaxException e) {
-      throw new RuntimeException("outputDirURI is not valid - '" + _spec.getOutputDirURI() + "'");
-    }
-    PinotFS outputDirFS = PinotFSFactory.create(outputDirURI.getScheme());
-
-    //Get list of files to process
-    String[] files;
-    try {
-      files = outputDirFS.listFiles(outputDirURI, true);
-    } catch (IOException e) {
-      throw new RuntimeException("Unable to list all files under outputDirURI - '" + outputDirURI + "'");
-    }
-    List<String> segmentUris = new ArrayList<>();
-    for (String file : files) {
-      URI uri = URI.create(file);
-      if (uri.getPath().endsWith(Constants.TAR_GZ_FILE_EXT)) {
-        URI updatedURI = SegmentPushUtils
-            .generateSegmentTarURI(outputDirURI, uri, _spec.getPushJobSpec().getSegmentUriPrefix(),
-                _spec.getPushJobSpec().getSegmentUriSuffix());
-        segmentUris.add(updatedURI.toString());
-      }
-    }
-    try {
-      SegmentPushUtils.sendSegmentUris(_spec, segmentUris);
-    } catch (RetriableOperationException | AttemptsExceededException e) {
-      throw new RuntimeException(e);
-    }
+  public void uploadSegments(Map<String, String> segmentsUriToTarPathMap)
+      throws AttemptsExceededException, RetriableOperationException {
+    SegmentPushUtils.sendSegmentUris(_spec, new ArrayList<>(segmentsUriToTarPathMap.keySet()));

Review Comment:
   same 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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r964298816


##########
pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/BaseSegmentPushJobRunner.java:
##########
@@ -0,0 +1,146 @@
+/**
+ * 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.plugin.ingestion.batch.common;
+
+import java.io.File;
+import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.segment.local.utils.ConsistentDataPushUtils;
+import org.apache.pinot.segment.local.utils.SegmentPushUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.filesystem.PinotFS;
+import org.apache.pinot.spi.filesystem.PinotFSFactory;
+import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+
+public abstract class BaseSegmentPushJobRunner implements IngestionJobRunner {
+
+  protected SegmentGenerationJobSpec _spec;
+  protected String[] _files;
+  protected PinotFS _outputDirFS;
+  protected URI _outputDirURI;
+  protected TableConfig _tableConfig;
+  protected boolean _consistentPushEnabled;
+
+  /**
+   * Initialize BaseSegmentPushJobRunner with SegmentGenerationJobSpec
+   * Checks for required parameters in the spec and enablement of consistent data push.
+   */
+  @Override
+  public void init(SegmentGenerationJobSpec spec) {
+    _spec = spec;
+    if (_spec.getPushJobSpec() == null) {
+      throw new RuntimeException("Missing PushJobSpec");
+    }
+
+    // Read Table spec
+    if (_spec.getTableSpec() == null) {
+      throw new RuntimeException("Missing tableSpec");
+    }
+
+    // Read Table config
+    if (_spec.getTableSpec().getTableConfigURI() == null) {
+      throw new RuntimeException("Missing property 'tableConfigURI' in 'tableSpec'");
+    }
+
+    _tableConfig = SegmentGenerationUtils.getTableConfig(_spec.getTableSpec().getTableConfigURI(), spec.getAuthToken());
+

Review Comment:
   done



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r964340776


##########
pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/BaseSegmentPushJobRunner.java:
##########
@@ -0,0 +1,147 @@
+/**
+ * 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.plugin.ingestion.batch.common;
+
+import java.io.File;
+import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.segment.local.utils.ConsistentDataPushUtils;
+import org.apache.pinot.segment.local.utils.SegmentPushUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.filesystem.PinotFS;
+import org.apache.pinot.spi.filesystem.PinotFSFactory;
+import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+
+public abstract class BaseSegmentPushJobRunner implements IngestionJobRunner {
+
+  protected SegmentGenerationJobSpec _spec;
+  protected String[] _files;
+  protected PinotFS _outputDirFS;
+  protected URI _outputDirURI;
+  protected TableConfig _tableConfig;
+  protected boolean _consistentPushEnabled;
+
+  /**
+   * Initialize BaseSegmentPushJobRunner with SegmentGenerationJobSpec
+   * Checks for required parameters in the spec and enablement of consistent data push.
+   */
+  @Override
+  public void init(SegmentGenerationJobSpec spec) {
+    _spec = spec;
+    if (_spec.getPushJobSpec() == null) {
+      throw new RuntimeException("Missing PushJobSpec");
+    }
+
+    // Read Table spec
+    if (_spec.getTableSpec() == null) {
+      throw new RuntimeException("Missing tableSpec");
+    }
+
+    // Read Table config
+    if (_spec.getTableSpec().getTableConfigURI() == null) {
+      throw new RuntimeException("Missing property 'tableConfigURI' in 'tableSpec'");
+    }
+
+    _tableConfig = SegmentGenerationUtils.getTableConfig(_spec.getTableSpec().getTableConfigURI(), spec.getAuthToken());
+    _consistentPushEnabled = ConsistentDataPushUtils.consistentDataPushEnabled(_tableConfig);
+  }
+
+  /**
+   * Initialize filesystems and obtain the raw input files for upload.
+   */
+  public void initFileSys() {
+    // Init all file systems
+    List<PinotFSSpec> pinotFSSpecs = _spec.getPinotFSSpecs();
+    for (PinotFSSpec pinotFSSpec : pinotFSSpecs) {
+      PinotFSFactory.register(pinotFSSpec.getScheme(), pinotFSSpec.getClassName(), new PinotConfiguration(pinotFSSpec));
+    }
+
+    // Get outputFS for writing output Pinot segments
+    try {
+      _outputDirURI = new URI(_spec.getOutputDirURI());
+      if (_outputDirURI.getScheme() == null) {
+        _outputDirURI = new File(_spec.getOutputDirURI()).toURI();
+      }
+    } catch (URISyntaxException e) {
+      throw new RuntimeException("outputDirURI is not valid - '" + _spec.getOutputDirURI() + "'");
+    }
+    _outputDirFS = PinotFSFactory.create(_outputDirURI.getScheme());
+
+    // Get list of files to process
+    try {
+      _files = _outputDirFS.listFiles(_outputDirURI, true);
+    } catch (IOException e) {
+      throw new RuntimeException("Unable to list all files under outputDirURI - '" + _outputDirURI + "'");
+    }
+  }
+
+  /**
+   * Returns segment names, which will be supplied to the segment replacement protocol as the new set of segments to
+   * atomically update when consistent data push is enabled.
+   * @param segmentsUriToTarPathMap Map from segment URI to corresponding tar path. Either the URIs (keys), the
+   *                                tarPaths (values), or both may be used depending on upload mode.
+   */
+  public List<String> getSegmentsToReplace(Map<String, String> segmentsUriToTarPathMap) {
+    return SegmentPushUtils.getSegmentNames(segmentsUriToTarPathMap.values());

Review Comment:
   Let's put the logic here instead of putting in `Util` class. We usually extract the logic to `Util` class when the same code needs to be called in multiple places to avoid the duplicate code. In this case, we only need to use this function 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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r964080138


##########
pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/ControllerRequestURLBuilder.java:
##########
@@ -370,6 +370,12 @@ public String forDeleteTableWithType(String tableName, String tableType) {
     return StringUtil.join("/", _baseUrl, "tables", tableName + "?type=" + tableType);
   }
 
+  public String forSegmentListAPIWithTableTypeAndExcludeReplacedSegments(String tableName, String tableType,

Review Comment:
   Thanks - unified all under `forSegmentListAPI` and introduced overloaded methods for default parameters. 



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] snleee commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
snleee commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r963952324


##########
pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/BaseSegmentPushJobRunner.java:
##########
@@ -0,0 +1,140 @@
+/**
+ * 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.plugin.ingestion.batch.common;
+
+import java.io.File;
+import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.segment.local.utils.ConsistentDataPushUtils;
+import org.apache.pinot.segment.local.utils.SegmentPushUtils;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.filesystem.PinotFS;
+import org.apache.pinot.spi.filesystem.PinotFSFactory;
+import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+
+public abstract class BaseSegmentPushJobRunner implements IngestionJobRunner {
+
+  protected SegmentGenerationJobSpec _spec;
+  protected String[] _files;
+  protected PinotFS _outputDirFS;
+  protected URI _outputDirURI;
+  protected boolean _consistentPushEnabled;
+
+  /**
+   * Initialize BaseSegmentPushJobRunner with SegmentGenerationJobSpec
+   * Checks for required parameters in the spec and enablement of consistent data push.
+   */
+  @Override
+  public void init(SegmentGenerationJobSpec spec) {
+    _spec = spec;
+    if (_spec.getPushJobSpec() == null) {
+      throw new RuntimeException("Missing PushJobSpec");
+    }
+
+    // Read Table spec
+    if (_spec.getTableSpec() == null) {
+      throw new RuntimeException("Missing tableSpec");
+    }
+
+    // Read Table config
+    if (_spec.getTableSpec().getTableConfigURI() == null) {
+      throw new RuntimeException("Missing property 'tableConfigURI' in 'tableSpec'");
+    }
+
+    _consistentPushEnabled = ConsistentDataPushUtils.consistentDataPushEnabled(_spec);
+  }
+
+  /**
+   * Initialize filesystems and obtain the raw input files for upload.
+   */
+  public void initFileSys() {
+    // Init all file systems
+    List<PinotFSSpec> pinotFSSpecs = _spec.getPinotFSSpecs();
+    for (PinotFSSpec pinotFSSpec : pinotFSSpecs) {
+      PinotFSFactory.register(pinotFSSpec.getScheme(), pinotFSSpec.getClassName(), new PinotConfiguration(pinotFSSpec));
+    }
+
+    // Get outputFS for writing output Pinot segments
+    try {
+      _outputDirURI = new URI(_spec.getOutputDirURI());
+      if (_outputDirURI.getScheme() == null) {
+        _outputDirURI = new File(_spec.getOutputDirURI()).toURI();
+      }
+    } catch (URISyntaxException e) {
+      throw new RuntimeException("outputDirURI is not valid - '" + _spec.getOutputDirURI() + "'");
+    }
+    _outputDirFS = PinotFSFactory.create(_outputDirURI.getScheme());
+
+    // Get list of files to process
+    try {
+      _files = _outputDirFS.listFiles(_outputDirURI, true);
+    } catch (IOException e) {
+      throw new RuntimeException("Unable to list all files under outputDirURI - '" + _outputDirURI + "'");
+    }
+  }
+
+  /**
+   * Returns segment names, which will be supplied to the segment replacement protocol as the new set of segments to
+   * atomically update when consistent data push is enabled.
+   * @param segmentsUriToTarPathMap Map from segment URI to corresponding tar path. Either the URIs (keys), the
+   *                                tarPaths (values), or both may be used depending on upload mode.
+   */
+  public abstract List<String> getSegmentsToReplace(Map<String, String> segmentsUriToTarPathMap);
+
+  /**
+   * Upload segments supplied in segmentsUriToTarPathMap.
+   * @param segmentsUriToTarPathMap Map from segment URI to corresponding tar path. Either the URIs (keys), the
+   *                                tarPaths (values), or both may be used depending on upload mode.
+   */
+  public abstract void uploadSegments(Map<String, String> segmentsUriToTarPathMap)
+      throws Exception;
+
+  /**
+   * Runs the main logic of the segment push job runner.
+   * First initialize the filesystem, then upload the segments, while optionally configured to be wrapped around by
+   * the consistent data push protocol.
+   */
+  @Override
+  public void run() {
+    initFileSys();
+    Map<URI, String> uriToLineageEntryIdMap = null;
+    try {
+      Map<String, String> segmentsUriToTarPathMap =
+          SegmentPushUtils.getSegmentUriToTarPathMap(_outputDirURI, _spec.getPushJobSpec(), _files);
+      if (_consistentPushEnabled) {
+        List<String> segmentsToReplace = getSegmentsToReplace(segmentsUriToTarPathMap);
+        uriToLineageEntryIdMap = ConsistentDataPushUtils.preUpload(_spec, segmentsToReplace);
+      }
+      uploadSegments(segmentsUriToTarPathMap);
+      if (_consistentPushEnabled) {
+        ConsistentDataPushUtils.postUpload(_spec, uriToLineageEntryIdMap);
+      }
+    } catch (Exception e) {
+      if (_consistentPushEnabled) {
+        ConsistentDataPushUtils.handleUploadException(_spec, uriToLineageEntryIdMap, e);
+      }
+      throw new RuntimeException(e);

Review Comment:
   We can throw `e` here



##########
pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/ControllerRequestURLBuilder.java:
##########
@@ -370,6 +370,12 @@ public String forDeleteTableWithType(String tableName, String tableType) {
     return StringUtil.join("/", _baseUrl, "tables", tableName + "?type=" + tableType);
   }
 
+  public String forSegmentListAPIWithTableTypeAndExcludeReplacedSegments(String tableName, String tableType,

Review Comment:
   Ideally, I think that we should have only 1 API for `forSegmentListAPI` with all parameters. I currently see that there are 2 (`forSegmentListAPI`, `forSegmentListAPIWithTableType`) and we are adding another API here. Can we clean this up?



##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java:
##########
@@ -0,0 +1,285 @@
+/**
+ * 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.segment.local.utils;
+
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.common.auth.AuthProviderUtils;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.apache.pinot.spi.utils.retry.RetryPolicy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class ConsistentDataPushUtils {
+  private ConsistentDataPushUtils() {
+  }
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentPushUtils.class);
+  private static final FileUploadDownloadClient FILE_UPLOAD_DOWNLOAD_CLIENT = new FileUploadDownloadClient();
+  private static final RetryPolicy DEFAULT_RETRY_POLICY = RetryPolicies.exponentialBackoffRetryPolicy(5, 10_000L, 2.0);
+  public static final String SEGMENT_NAME_POSTFIX = "segment.name.postfix";
+
+  /**
+   * Checks for enablement of consistent data push. If enabled, fetch the list of segments to be replaced, then
+   * invoke startReplaceSegments API and returns a map of controller URI to lineage entry IDs.
+   * If not, returns an empty hashmap.
+   */
+  public static Map<URI, String> preUpload(SegmentGenerationJobSpec spec, List<String> segmentsTo)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    boolean consistentDataPushEnabled = consistentDataPushEnabled(spec);

Review Comment:
   We already check if the consistent push is enabled in the caller side. We can remove this.



##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java:
##########
@@ -0,0 +1,285 @@
+/**
+ * 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.segment.local.utils;
+
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.common.auth.AuthProviderUtils;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest;
+import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.retry.RetryPolicies;
+import org.apache.pinot.spi.utils.retry.RetryPolicy;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class ConsistentDataPushUtils {
+  private ConsistentDataPushUtils() {
+  }
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SegmentPushUtils.class);
+  private static final FileUploadDownloadClient FILE_UPLOAD_DOWNLOAD_CLIENT = new FileUploadDownloadClient();
+  private static final RetryPolicy DEFAULT_RETRY_POLICY = RetryPolicies.exponentialBackoffRetryPolicy(5, 10_000L, 2.0);
+  public static final String SEGMENT_NAME_POSTFIX = "segment.name.postfix";
+
+  /**
+   * Checks for enablement of consistent data push. If enabled, fetch the list of segments to be replaced, then
+   * invoke startReplaceSegments API and returns a map of controller URI to lineage entry IDs.
+   * If not, returns an empty hashmap.
+   */
+  public static Map<URI, String> preUpload(SegmentGenerationJobSpec spec, List<String> segmentsTo)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    boolean consistentDataPushEnabled = consistentDataPushEnabled(spec);
+    LOGGER.info("Consistent data push is: {}", consistentDataPushEnabled ? "enabled" : "disabled");
+    Map<URI, String> uriToLineageEntryIdMap = null;
+    if (consistentDataPushEnabled) {
+      LOGGER.info("Start consistent push for table: " + rawTableName);
+      Map<URI, List<String>> uriToExistingOfflineSegments = getSegmentsToReplace(spec, rawTableName);
+      LOGGER.info("Existing segments for table {}: " + uriToExistingOfflineSegments, rawTableName);
+      LOGGER.info("New segments for table: {}: " + segmentsTo, rawTableName);
+      uriToLineageEntryIdMap = startReplaceSegments(spec, uriToExistingOfflineSegments, segmentsTo);
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * uriToLineageEntryIdMap is non-empty if and only if consistent data push is enabled.
+   * If uriToLineageEntryIdMap is non-empty, end the consistent data push protocol for each controller.
+   */
+  public static void postUpload(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap)
+      throws Exception {
+    String rawTableName = spec.getTableSpec().getTableName();
+    if (uriToLineageEntryIdMap != null && !uriToLineageEntryIdMap.isEmpty()) {
+      LOGGER.info("End consistent push for table: " + rawTableName);
+      endReplaceSegments(spec, uriToLineageEntryIdMap);
+    }
+  }
+
+  /**
+   * Builds a map of controller URI to startReplaceSegments URI for each Pinot cluster in the spec.
+   */
+  public static Map<URI, URI> getStartReplaceSegmentUris(SegmentGenerationJobSpec spec, String rawTableName) {
+    Map<URI, URI> baseUriToStartReplaceSegmentUriMap = new HashMap<>();
+    for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
+      URI controllerURI;
+      try {
+        controllerURI = new URI(pinotClusterSpec.getControllerURI());
+        baseUriToStartReplaceSegmentUriMap.put(controllerURI,
+            FileUploadDownloadClient.getStartReplaceSegmentsURI(controllerURI, rawTableName,
+                TableType.OFFLINE.toString(), true));
+      } catch (URISyntaxException e) {
+        throw new RuntimeException("Got invalid controller uri - '" + pinotClusterSpec.getControllerURI() + "'");
+      }
+    }
+    return baseUriToStartReplaceSegmentUriMap;
+  }
+
+  /**
+   * Starts consistent data push protocol for each Pinot cluster in the spec.
+   * Returns a map of controller URI to segment lineage entry ID.
+   */
+  public static Map<URI, String> startReplaceSegments(SegmentGenerationJobSpec spec,
+      Map<URI, List<String>> uriToSegmentsFrom, List<String> segmentsTo)
+      throws Exception {
+    Map<URI, String> uriToLineageEntryIdMap = new HashMap<>();
+    String rawTableName = spec.getTableSpec().getTableName();
+    Map<URI, URI> segmentsUris = getStartReplaceSegmentUris(spec, rawTableName);
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    LOGGER.info("Start replace segment URIs: " + segmentsUris);
+
+    for (Map.Entry<URI, URI> entry : segmentsUris.entrySet()) {
+      URI controllerUri = entry.getKey();
+      URI startSegmentUri = entry.getValue();
+      List<String> segmentsFrom = uriToSegmentsFrom.get(controllerUri);
+
+      StartReplaceSegmentsRequest startReplaceSegmentsRequest =
+          new StartReplaceSegmentsRequest(segmentsFrom, segmentsTo);
+      DEFAULT_RETRY_POLICY.attempt(() -> {
+        try {
+          SimpleHttpResponse response =
+              FILE_UPLOAD_DOWNLOAD_CLIENT.startReplaceSegments(startSegmentUri, startReplaceSegmentsRequest,
+                  authProvider);
+          String responseString = response.getResponse();
+          LOGGER.info(
+              "Got response {}: {} while sending start replace segment request for table: {}, uploadURI: {}, request:"
+                  + " {}", response.getStatusCode(), responseString, rawTableName, startSegmentUri,
+              startReplaceSegmentsRequest);
+          String segmentLineageEntryId =
+              JsonUtils.stringToJsonNode(responseString).get("segmentLineageEntryId").asText();
+          uriToLineageEntryIdMap.put(controllerUri, segmentLineageEntryId);
+          return true;
+        } catch (SocketTimeoutException se) {
+          // In case of the timeout, we should re-try.
+          return false;
+        } catch (HttpErrorStatusException e) {
+          if (e.getStatusCode() >= 500) {
+            return false;
+          } else {
+            if (e.getStatusCode() == Response.Status.NOT_FOUND.getStatusCode()) {
+              LOGGER.error("Table: {} not found when sending request: {}", rawTableName, startSegmentUri);
+            }
+            throw e;
+          }
+        }
+      });
+    }
+    return uriToLineageEntryIdMap;
+  }
+
+  /**
+   * Ends consistent data push protocol for each Pinot cluster in the spec.
+   */
+  public static void endReplaceSegments(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap)
+      throws Exception {
+    AuthProvider authProvider = AuthProviderUtils.makeAuthProvider(spec.getAuthToken());
+    String rawTableName = spec.getTableSpec().getTableName();
+    for (URI controllerUri : uriToLineageEntryIdMap.keySet()) {
+      String segmentLineageEntryId = uriToLineageEntryIdMap.get(controllerUri);
+      URI uri =
+          FileUploadDownloadClient.getEndReplaceSegmentsURI(controllerUri, rawTableName, TableType.OFFLINE.toString(),
+              segmentLineageEntryId);
+      DEFAULT_RETRY_POLICY.attempt(() -> {
+        try {
+          SimpleHttpResponse response =
+              FILE_UPLOAD_DOWNLOAD_CLIENT.endReplaceSegments(uri, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS, authProvider);
+          LOGGER.info("Got response {}: {} while sending end replace segment request for table: {}, uploadURI: {}",
+              response.getStatusCode(), response.getResponse(), rawTableName, uri);
+          return true;
+        } catch (SocketTimeoutException se) {
+          // In case of the timeout, we should re-try.
+          return false;
+        } catch (HttpErrorStatusException e) {
+          if (e.getStatusCode() >= 500) {
+            return false;
+          } else {
+            throw e;
+          }
+        }
+      });
+    }
+  }
+
+  /**
+   * Revert segment lineage entry when exception gets caught. This revert request is called at best effort.
+   * If the revert call fails at this point, the next startReplaceSegment call will do the cleanup
+   * by marking the previous entry to "REVERTED" and cleaning up the leftover segments.
+   */
+  public static void handleUploadException(SegmentGenerationJobSpec spec, Map<URI, String> uriToLineageEntryIdMap,
+      Exception exception) {
+    if (uriToLineageEntryIdMap != null) {
+      LOGGER.error("Exception when pushing segments. Marking segment lineage entry to 'REVERTED'.", exception);
+      String rawTableName = spec.getTableSpec().getTableName();
+      for (Map.Entry<URI, String> entry : uriToLineageEntryIdMap.entrySet()) {
+        String segmentLineageEntryId = entry.getValue();
+        try {
+          URI uri = FileUploadDownloadClient.getRevertReplaceSegmentsURI(entry.getKey(), rawTableName,
+              TableType.OFFLINE.name(), segmentLineageEntryId, true);
+          SimpleHttpResponse response = FILE_UPLOAD_DOWNLOAD_CLIENT.revertReplaceSegments(uri);
+          LOGGER.info("Got response {}: {} while sending revert replace segment request for table: {}, uploadURI: {}",
+              response.getStatusCode(), response.getResponse(), rawTableName, entry.getKey());
+        } catch (URISyntaxException | HttpErrorStatusException | IOException e) {
+          LOGGER.error("Exception when sending revert replace segment request to controller: {} for table: {}",
+              entry.getKey(), rawTableName, e);
+        }
+      }
+    }
+  }
+
+  public static TableConfig getTableConfig(SegmentGenerationJobSpec spec) {
+    String rawTableName = spec.getTableSpec().getTableName();
+    TableConfig tableConfig =
+        SegmentGenerationUtils.getTableConfig(spec.getTableSpec().getTableConfigURI(), spec.getAuthToken());
+    Preconditions.checkState(tableConfig != null, "Failed to find table config for table: %s", rawTableName);
+    return tableConfig;
+  }
+
+  public static boolean consistentDataPushEnabled(SegmentGenerationJobSpec spec) {
+    TableConfig tableConfig = getTableConfig(spec);
+    // Enable consistent data push only if "consistentDataPush" is set to true in batch ingestion config and the
+    // table is REFRESH use case.
+    // TODO: Remove the check for REFRESH when we support consistent push for APPEND table
+    return "REFRESH".equalsIgnoreCase(IngestionConfigUtils.getBatchSegmentIngestionType(tableConfig))
+        && IngestionConfigUtils.getBatchSegmentIngestionConsistentDataPushEnabled(tableConfig);
+  }
+
+  /**
+   * Returns a map of controller URI to a list of existing OFFLINE segments.
+   */
+  public static Map<URI, List<String>> getSegmentsToReplace(SegmentGenerationJobSpec spec, String rawTableName) {
+    Map<URI, List<String>> uriToOfflineSegments = new HashMap<>();
+    for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
+      URI controllerURI;
+      List<String> offlineSegments;
+      try {
+        controllerURI = new URI(pinotClusterSpec.getControllerURI());
+        Map<String, List<String>> segments =
+            FILE_UPLOAD_DOWNLOAD_CLIENT.getSegments(controllerURI, rawTableName, TableType.OFFLINE, true);
+        offlineSegments = segments.get(TableType.OFFLINE.toString());
+        uriToOfflineSegments.put(controllerURI, offlineSegments);
+      } catch (URISyntaxException e) {
+        throw new RuntimeException("Got invalid controller uri - '" + pinotClusterSpec.getControllerURI() + "'");
+      } catch (Exception e) {
+        e.printStackTrace();

Review Comment:
   We should avoid to use `e.printStrackTrace()`. Can you throw the exception here and let the caller handle the exception?



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r958994759


##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentUploadIntegrationTest.java:
##########
@@ -81,10 +89,15 @@ protected List<String> getBloomFilterColumns() {
     return null;
   }
 
+  @BeforeMethod
+  public void setUpTest()

Review Comment:
   Calling `TestUtils.ensureDirectoriesExistAndEmpty(_tempDir, _segmentDir, _tarDir)` in `setUp` is insufficient as there is another test case `testUploadAndQuery` that leaves data in the tmp directories.
   
   Using two table names will also lead to more code for now, since it seems like `BaseClusterIntegrationTest` is written in a way that assumes one table/config per subclass, e.g. `getTableName`, `createOfflineTableConfig`. 
   
   Hence, I prefer to keep the test structure as is with the before and after method annotations. Another way is to create a different test class altogether, but combining the two here to not prevent slow down from too many tests. 



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org


[GitHub] [pinot] yuanbenson commented on a diff in pull request #9295: Enable Consistent Data Push for Standalone Segment Push Job Runners

Posted by GitBox <gi...@apache.org>.
yuanbenson commented on code in PR #9295:
URL: https://github.com/apache/pinot/pull/9295#discussion_r959048302


##########
pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java:
##########
@@ -747,6 +751,54 @@ public SimpleHttpResponse uploadSegment(URI uri, String segmentName, InputStream
     return uploadSegment(uri, segmentName, inputStream, null, parameters, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
   }
 
+  /**
+   * Returns a map from a given tableType to a list of segments for that given tableType (OFFLINE or REALTIME)
+   * If tableType is left unspecified, both OFFLINE and REALTIME segments will be returned in the map.
+   */
+  public Map<String, List<String>> getSegments(URI uri, String rawTableName, @Nullable String tableType,
+      boolean excludeReplacedSegments)
+      throws URISyntaxException, IOException {
+    List<String> tableTypes;
+    if (tableType == null || tableType.isEmpty()) {
+      tableTypes = Arrays.asList(TableType.OFFLINE.toString(), TableType.REALTIME.toString());
+    } else {
+      tableTypes = Arrays.asList(tableType);
+    }
+    ControllerRequestURLBuilder controllerRequestURLBuilder = ControllerRequestURLBuilder.baseUrl(uri.toString());
+    Map<String, List<String>> tableTypeToSegments = new HashMap<>();
+    for (String tableTypeToFilter : tableTypes) {
+      List<String> segments = new ArrayList<>();
+      RequestBuilder requestBuilder = RequestBuilder.get(
+          controllerRequestURLBuilder.forSegmentListAPIWithTableTypeAndExcludeReplacedSegments(rawTableName,
+              tableTypeToFilter, excludeReplacedSegments)).setVersion(HttpVersion.HTTP_1_1);
+      HttpClient.setTimeout(requestBuilder, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
+      SimpleHttpResponse response;
+      try {
+        response = HttpClient.wrapAndThrowHttpException(_httpClient.sendRequest(requestBuilder.build()));
+      } catch (HttpErrorStatusException e) {
+        tableTypeToSegments.put(tableTypeToFilter, segments);
+        continue;
+      }
+      String responseString = response.getResponse();
+      JsonNode responseJsonNode = JsonUtils.stringToJsonNode(responseString);
+      Iterator<JsonNode> responseElements = responseJsonNode.elements();
+      while (responseElements.hasNext()) {

Review Comment:
   Could you clarify a bit more on the input parameters/ask for this method we are talking about? I'll be glad to extract that out if it does not already exist. 
   
   Ohh I see, I'll look into possibilities of reusing `getSelectedSegments` API that you previously brought in. That may be used when we support APPEND tables anyways! Thx. 



-- 
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: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org