#12117 Support for Server & Controller API to check for Segments reload of a table in servers

pinot-common/src/main/java/org/apache/pinot/common/restlet/resources/SegmentsReloadCheckResponse.java

@@ -0,0 +1,48 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.common.restlet.resources;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+
+/**
+ * This class gives the data of a server if there exists any segments that need to be reloaded
+ *
+ * It has details of server id and returns true/false if there are any segments to be reloaded or not.
+ */
+public class SegmentsReloadCheckResponse {
+  private final boolean _needReload;
+  private final String _serverInstanceId;
+
+  public String getServerInstanceId() {
+    return _serverInstanceId;
+  }
+
+  public boolean getReload() {
+    return _needReload;
+  }
+
+  @JsonCreator
+  public SegmentsReloadCheckResponse(@JsonProperty("needReload") boolean needReload,
+      @JsonProperty("serverInstanceId") String serverInstanceId) {
+    _needReload = needReload;
+    _serverInstanceId = serverInstanceId;
+  }
+}

pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java

@@ -1013,6 +1013,22 @@ public boolean tryLoadExistingSegment(SegmentZKMetadata zkMetadata, IndexLoading
     }
   }
 
+  boolean needReloadSegment(SegmentZKMetadata zkMetadata, IndexLoadingConfig indexLoadingConfig)
+      throws Exception {
+    String segmentName = zkMetadata.getSegmentName();
+    SegmentDirectory segmentDirectory =
+        tryInitSegmentDirectory(segmentName, String.valueOf(zkMetadata.getCrc()), indexLoadingConfig);
+    try {
+      Schema schema = indexLoadingConfig.getSchema();
+      //if re processing or reload is needed on a segment then return true
+      return ImmutableSegmentLoader.needPreprocess(segmentDirectory, indexLoadingConfig, schema);
+    } catch (Exception e) {
+      _logger.error("Failed to check if reload is needed for a segment {}", segmentName, e);
+      closeSegmentDirectoryQuietly(segmentDirectory);
+      throw new Exception(String.format("Failed to perform reload check on the segment %s", segmentName));
+    }
+  }
+
   @Nullable
   private SegmentDirectory tryInitSegmentDirectory(String segmentName, String segmentCrc,
       IndexLoadingConfig indexLoadingConfig) {
@@ -1024,6 +1040,28 @@ private SegmentDirectory tryInitSegmentDirectory(String segmentName, String segm
     }
   }
 
+  @Override
+  public boolean needReloadSegments()
+      throws Exception {
+    IndexLoadingConfig indexLoadingConfig = fetchIndexLoadingConfig();
+    List<SegmentDataManager> segmentDataManagers = acquireAllSegments();
+    boolean needReload = false;
+    try {
+      for (SegmentDataManager segmentDataManager : segmentDataManagers) {
+        SegmentZKMetadata segmentZKMetadata = fetchZKMetadata(segmentDataManager.getSegmentName());
+        if (needReloadSegment(segmentZKMetadata, indexLoadingConfig)) {
+          needReload = true;
+          break;
+        }
+      }
+    } finally {
+      for (SegmentDataManager segmentDataManager : segmentDataManagers) {
+        releaseSegment(segmentDataManager);
+      }
+    }
+    return needReload;
+  }
+
   private SegmentDirectory initSegmentDirectory(String segmentName, String segmentCrc,
       IndexLoadingConfig indexLoadingConfig)
       throws Exception {

pinot-core/src/test/java/org/apache/pinot/core/data/manager/BaseTableDataManagerTest.java

@@ -49,6 +49,7 @@
 import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
 import org.apache.pinot.segment.spi.creator.SegmentVersion;
 import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl;
+import org.apache.pinot.segment.spi.store.SegmentDirectory;
 import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths;
 import org.apache.pinot.spi.config.instance.InstanceDataManagerConfig;
 import org.apache.pinot.spi.config.table.TableConfig;
@@ -289,6 +290,30 @@ public void testReloadSegmentAddIndex()
     assertTrue(hasInvertedIndex(indexDir, LONG_COLUMN));
   }
 
+  @Test
+  public void testNeedReloadSegmentAddIndex()
+      throws Exception {
+    File indexDir = createSegment(SegmentVersion.v3, 5);
+    long crc = getCRC(indexDir);
+    assertFalse(hasInvertedIndex(indexDir, STRING_COLUMN));
+    assertFalse(hasInvertedIndex(indexDir, LONG_COLUMN));
+
+    // Same CRCs so load the local segment directory directly.
+    SegmentZKMetadata zkMetadata = mock(SegmentZKMetadata.class);
+    when(zkMetadata.getCrc()).thenReturn(crc);
+    when(zkMetadata.getSegmentName()).thenReturn(SEGMENT_NAME);
+    SegmentMetadata localMetadata = mock(SegmentMetadata.class);
+    when(localMetadata.getCrc()).thenReturn(Long.toString(crc));
+    SegmentDirectory segmentDirectory = mock(SegmentDirectory.class);
+    // Require to add indices.
+    IndexLoadingConfig indexLoadingConfig = new IndexLoadingConfig();
+    indexLoadingConfig.setInvertedIndexColumns(new HashSet<>(Arrays.asList(STRING_COLUMN, LONG_COLUMN)));
+    BaseTableDataManager tableDataManager = createTableManager();
+    assertTrue(tableDataManager.needReloadSegment(zkMetadata, indexLoadingConfig));
+    tableDataManager.reloadSegment(SEGMENT_NAME, indexLoadingConfig, zkMetadata, localMetadata, null, false);
+    assertFalse(tableDataManager.needReloadSegment(zkMetadata, indexLoadingConfig));
+  }
+
   @Test
   public void testReloadSegmentForceDownload()
       throws Exception {

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/TableDataManager.java

@@ -122,6 +122,13 @@ void addNewOnlineSegment(SegmentZKMetadata zkMetadata, IndexLoadingConfig indexL
    */
   boolean tryLoadExistingSegment(SegmentZKMetadata zkMetadata, IndexLoadingConfig indexLoadingConfig);
 
+  /**
+   * Check if reload is needed for any of the segments of a table
+   * @return true if reload is needed for any of the segments and false otherwise
+   */
+  boolean needReloadSegments()
+      throws Exception;
+
   /**
    * Downloads a segment and loads it into the table.
    * NOTE: This method is part of the implementation detail of {@link #addOnlineSegment(String)}.

pinot-server/src/main/java/org/apache/pinot/server/api/resources/TablesResource.java

@@ -64,6 +64,7 @@
 import org.apache.pinot.common.response.server.TableIndexMetadataResponse;
 import org.apache.pinot.common.restlet.resources.ResourceUtils;
 import org.apache.pinot.common.restlet.resources.SegmentConsumerInfo;
+import org.apache.pinot.common.restlet.resources.SegmentsReloadCheckResponse;
 import org.apache.pinot.common.restlet.resources.TableMetadataInfo;
 import org.apache.pinot.common.restlet.resources.TableSegmentValidationInfo;
 import org.apache.pinot.common.restlet.resources.TableSegments;
@@ -952,4 +953,28 @@ public TableSegmentValidationInfo validateTableSegmentState(
     }
     return new TableSegmentValidationInfo(true, maxEndTimeMs);
   }
+
+  @GET
+  @Path("/tables/{tableName}/segments/needReload")
+  @Produces(MediaType.APPLICATION_JSON)
+  @ApiOperation(value = "Checks if reload is needed on any segment", notes = "Returns true if reload is required on"
+      + " any segment in this server")
+  @ApiResponses(value = {
+      @ApiResponse(code = 200, message = "Success", response = TableSegments.class), @ApiResponse(code = 500,
+      message = "Internal Server error", response = ErrorInfo.class)
+  })
+  public String checkSegmentsReload(
+      @ApiParam(value = "Table Name with type", required = true) @PathParam("tableName") String tableName,
+      @Context HttpHeaders headers) {
+    tableName = DatabaseUtils.translateTableName(tableName, headers);
+    TableDataManager tableDataManager = ServerResourceUtils.checkGetTableDataManager(_serverInstance, tableName);
+    boolean needReload = false;
+    try {
+      needReload = tableDataManager.needReloadSegments();
+    } catch (Exception e) {
+      throw new WebApplicationException(e.getMessage(), Response.Status.INTERNAL_SERVER_ERROR);
+    }
+    return ResourceUtils.convertToJsonString(
+        new SegmentsReloadCheckResponse(needReload, tableDataManager.getInstanceId()));
+  }
 }