fix(pit): reduce overhead of PIT creation (#15412)

Author: RyanHolstien

metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ESGraphQueryDAO.java

@@ -99,6 +99,11 @@ SearchResponse executeSearch(@Nonnull SearchRequest searchRequest) {
     return delegate.executeSearch(searchRequest);
   }
 
+  @Override
+  public void cleanupPointInTime(String pitId) {
+    delegate.cleanupPointInTime(pitId);
+  }
+
   @Override
   public void destroy() throws Exception {
     // Shutdown the delegate if it's a GraphQueryPITDAO

metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ElasticSearchGraphService.java

@@ -60,6 +60,7 @@
 import lombok.extern.slf4j.Slf4j;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
+import org.opensearch.common.unit.TimeValue;
 import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
@@ -370,9 +371,22 @@ public RelatedEntitiesScrollResult scrollRelatedEntities(
     SearchHit[] searchHits = response.getHits().getHits();
     // Only return next scroll ID if there are more results, indicated by full size results
     String nextScrollId = null;
+    String pitId = response.pointInTimeId();
     if (searchHits.length == count && searchHits.length > 0) {
       Object[] sort = searchHits[searchHits.length - 1].getSortValues();
-      nextScrollId = new SearchAfterWrapper(sort, null, 0L).toScrollId();
+      if (pitId != null && keepAlive == null) {
+        throw new IllegalArgumentException("Should not set pitId without keepAlive");
+      }
+      long expirationTime =
+          keepAlive == null
+              ? 0L
+              : System.currentTimeMillis()
+                  + TimeValue.parseTimeValue(keepAlive, "keepAlive").millis();
+      nextScrollId = new SearchAfterWrapper(sort, pitId, expirationTime).toScrollId();
+    }
+    if (nextScrollId == null && pitId != null) {
+      // Last scroll, we clean up the pitId assuming user has gone through all data
+      graphReadDAO.cleanupPointInTime(pitId);
     }
 
     return RelatedEntitiesScrollResult.builder()

metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/GraphQueryBaseDAO.java

@@ -1664,4 +1664,9 @@ protected List<LineageRelationship> processSliceFutures(
 
     return allRelationships;
   }
+
+  @Override
+  public void cleanupPointInTime(String pitId) {
+    ESUtils.cleanupPointInTime(getClient(), pitId, "API Request");
+  }
 }

metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/GraphQueryDAO.java

@@ -47,4 +47,6 @@ SearchResponse getSearchResponse(
       @Nullable String scrollId,
       @Nullable String keepAlive,
       @Nullable Integer count);
+
+  void cleanupPointInTime(String pitId);
 }

metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/GraphQueryPITDAO.java

@@ -118,37 +118,53 @@ protected List<LineageRelationship> searchWithSlices(
       boolean allowPartialResults) {
 
     // Create slice-based search requests
-    List<CompletableFuture<List<LineageRelationship>>> sliceFutures = new ArrayList<>();
+    String pitId = null;
+    String keepAlive = config.getSearch().getGraph().getImpact().getKeepAlive();
+    try {
+      List<CompletableFuture<List<LineageRelationship>>> sliceFutures = new ArrayList<>();
+      pitId =
+          ESUtils.computePointInTime(
+              null,
+              keepAlive,
+              client,
+              opContext.getSearchContext().getIndexConvention().getIndexName(INDEX_NAME));
+      final String tempPitId = pitId;
 
-    for (int sliceId = 0; sliceId < slices; sliceId++) {
-      final int currentSliceId = sliceId;
+      for (int sliceId = 0; sliceId < slices; sliceId++) {
+        final int currentSliceId = sliceId;
 
-      CompletableFuture<List<LineageRelationship>> sliceFuture =
-          CompletableFuture.supplyAsync(
-              () -> {
-                return searchSingleSliceWithPit(
-                    opContext,
-                    query,
-                    lineageGraphFilters,
-                    visitedEntities,
-                    viaEntities,
-                    numHops,
-                    remainingHops,
-                    existingPaths,
-                    maxRelations,
-                    defaultPageSize,
-                    currentSliceId,
-                    slices,
-                    remainingTime,
-                    entityUrns,
-                    allowPartialResults);
-              },
-              pitExecutor); // Use dedicated thread pool with CallerRunsPolicy for backpressure
-      sliceFutures.add(sliceFuture);
-    }
+        CompletableFuture<List<LineageRelationship>> sliceFuture =
+            CompletableFuture.supplyAsync(
+                () -> {
+                  return searchSingleSliceWithPit(
+                      opContext,
+                      query,
+                      lineageGraphFilters,
+                      visitedEntities,
+                      viaEntities,
+                      numHops,
+                      remainingHops,
+                      existingPaths,
+                      maxRelations,
+                      defaultPageSize,
+                      currentSliceId,
+                      slices,
+                      remainingTime,
+                      entityUrns,
+                      allowPartialResults,
+                      tempPitId,
+                      keepAlive);
+                },
+                pitExecutor); // Use dedicated thread pool with CallerRunsPolicy for backpressure
+        sliceFutures.add(sliceFuture);
+      }
 
-    // Reuse the common slice coordination logic
-    return processSliceFutures(sliceFutures, remainingTime, allowPartialResults);
+      // Reuse the common slice coordination logic
+      return processSliceFutures(sliceFutures, remainingTime, allowPartialResults);
+    } finally {
+      // Clean up PIT to prevent hitting the limit
+      ESUtils.cleanupPointInTime(client, pitId, "lineage search: " + entityUrns);
+    }
   }
 
   /**
@@ -173,22 +189,14 @@ private List<LineageRelationship> searchSingleSliceWithPit(
       int totalSlices,
       long remainingTime,
       Set<Urn> entityUrns,
-      boolean allowPartialResults) {
+      boolean allowPartialResults,
+      String pitId,
+      String keepAlive) {
 
     List<LineageRelationship> sliceRelationships = new ArrayList<>();
-    String pitId = null;
     Object[] searchAfter = null;
-    String keepAlive = config.getSearch().getGraph().getImpact().getKeepAlive();
 
     try {
-      // Create initial PIT using existing utility method
-      pitId =
-          ESUtils.computePointInTime(
-              null,
-              keepAlive,
-              client,
-              opContext.getSearchContext().getIndexConvention().getIndexName(INDEX_NAME));
-
       // If maxRelations is -1 or 0, treat as unlimited (only bound by time)
       while (maxRelations <= 0 || sliceRelationships.size() < maxRelations) {
         // Check for thread interruption (from future.cancel(true))
@@ -293,9 +301,6 @@ private List<LineageRelationship> searchSingleSliceWithPit(
     } catch (Exception e) {
       log.error("Failed to execute PIT search for slice {}", sliceId, e);
       throw new RuntimeException("Failed to execute PIT search for slice " + sliceId, e);
-    } finally {
-      // Clean up PIT to prevent hitting the limit
-      ESUtils.cleanupPointInTime(client, pitId, "slice " + sliceId);
     }
 
     return sliceRelationships;

metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java

@@ -10,13 +10,15 @@
 import static com.linkedin.metadata.search.elasticsearch.query.request.SearchFieldConfig.KEYWORD_FIELDS;
 import static com.linkedin.metadata.search.elasticsearch.query.request.SearchFieldConfig.PATH_HIERARCHY_FIELDS;
 import static com.linkedin.metadata.utils.CriterionUtils.buildCriterion;
+import static org.opensearch.core.rest.RestStatus.TOO_MANY_REQUESTS;
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.google.common.collect.ImmutableList;
 import com.linkedin.data.schema.DataSchema;
 import com.linkedin.data.schema.MapDataSchema;
 import com.linkedin.data.schema.PathSpec;
 import com.linkedin.metadata.aspect.AspectRetriever;
+import com.linkedin.metadata.dao.throttle.APIThrottleException;
 import com.linkedin.metadata.models.EntitySpec;
 import com.linkedin.metadata.models.SearchableFieldSpec;
 import com.linkedin.metadata.models.StructuredPropertyUtils;
@@ -53,6 +55,7 @@
 import javax.annotation.Nullable;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.StringUtils;
+import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.search.CreatePitRequest;
 import org.opensearch.action.search.CreatePitResponse;
 import org.opensearch.action.search.DeletePitRequest;
@@ -1447,6 +1450,20 @@ public static QueryBuilder queryOptimize(QueryBuilder query, boolean considerSco
           new CreatePitRequest(TimeValue.parseTimeValue(keepAlive, "keepAlive"), false, indexArray);
       CreatePitResponse response = client.createPit(request, RequestOptions.DEFAULT);
       return response.getId();
+    } catch (OpenSearchStatusException ose) {
+      if (TOO_MANY_REQUESTS.equals(ose.status())) {
+        APIThrottleException throttleException =
+            new APIThrottleException(
+                TimeValue.parseTimeValue(keepAlive, "keepAlive").millis(),
+                "Too many point in times created, retry after keep alive has expired.");
+        try {
+          throttleException.initCause(ose);
+        } catch (IllegalStateException | IllegalArgumentException e) {
+          // Do nothing, can't fill in cause
+        }
+        throw throttleException;
+      }
+      throw ose;
     } catch (IOException e) {
       log.warn("Failed to generate PointInTime Identifier:", e);
       throw new IllegalStateException("Failed to generate PointInTime Identifier.", e);

metadata-io/src/test/java/com/linkedin/metadata/graph/elastic/ESGraphQueryDAOTest.java

@@ -1,8 +1,11 @@
 package com.linkedin.metadata.graph.elastic;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.isNull;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static org.testng.Assert.assertEquals;
@@ -24,6 +27,7 @@
 import com.linkedin.metadata.utils.metrics.MetricUtils;
 import io.datahubproject.metadata.context.OperationContext;
 import io.datahubproject.test.metadata.context.TestOperationContexts;
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -523,4 +527,20 @@ public void testDestroyWithNullDelegate() throws Exception {
 
     // Test passes if no exception is thrown
   }
+
+  @Test
+  public void testCleanupPointInTime() throws IOException {
+    SearchClientShim<?> mockClient = mock(SearchClientShim.class);
+
+    when(mockClient.getEngineType()).thenReturn(SearchClientShim.SearchEngineType.OPENSEARCH_2);
+    ESGraphQueryDAO dao =
+        new ESGraphQueryDAO(
+            mockClient, mockGraphServiceConfig, mockElasticSearchConfig, mockMetricUtils);
+
+    String pitId = "test-pit-id";
+    dao.cleanupPointInTime(pitId);
+
+    // Verify ESUtils.cleanupPointInTime was called with correct parameters
+    verify(mockClient, times(1)).deletePit(argThat(req -> req.getPitIds().contains(pitId)), any());
+  }
 }