@Override
  public void preloadSegments(IndexLoadingConfig indexLoadingConfig) {
    if (!_isPreloading) {
      return;
    }
    TableDataManager tableDataManager = _context.getTableDataManager();
    Preconditions.checkNotNull(tableDataManager, "Preloading segments requires tableDataManager");
    HelixManager helixManager = tableDataManager.getHelixManager();
    ExecutorService segmentPreloadExecutor = tableDataManager.getSegmentPreloadExecutor();
    // Preloading the segments with the snapshots of validDocIds for fast upsert metadata recovery.
    // Note that there is a waiting logic between the thread pool doing the segment preloading here and the
    // other helix threads about to process segment state transitions (e.g. taking segments from OFFLINE to ONLINE).
    // The thread doing the segment preloading here must complete before the other helix threads start to handle
    // segment state transitions. This is ensured by the lock here.
    _preloadLock.lock();
    try {
      // Check the flag again to ensure preloading happens only once.
      if (!_isPreloading) {
        return;
      }
      // From now on, the _isPreloading flag is true until the segments are preloaded.
      long startTime = System.currentTimeMillis();
      doPreloadSegments(tableDataManager, indexLoadingConfig, helixManager, segmentPreloadExecutor);
      long duration = System.currentTimeMillis() - startTime;
      _serverMetrics.addTimedTableValue(_tableNameWithType, ServerTimer.UPSERT_PRELOAD_TIME_MS, duration,
          TimeUnit.MILLISECONDS);
    } catch (Exception e) {
      // Even if preloading fails, we should continue to complete the initialization, so that TableDataManager can be
      // created. Once TableDataManager is created, no more segment preloading would happen, and the normal segment
      // loading logic would be used. The segments not being preloaded successfully here would be loaded via the
      // normal segment loading logic, the one doing more costly checks on the upsert metadata.
      _logger.warn("Failed to preload segments from partition: {} of table: {}, skipping", _partitionId,
          _tableNameWithType, e);
      _serverMetrics.addMeteredTableValue(_tableNameWithType, ServerMeter.UPSERT_PRELOAD_FAILURE, 1);
      if (e instanceof InterruptedException) {
        // Restore the interrupted status in case the upper callers want to check.
        Thread.currentThread().interrupt();
      }
    } finally {
      _isPreloading = false;
      _preloadLock.unlock();
    }
  }
@Test
  public void testPreloadSegments()
      throws Exception {
    String realtimeTableName = "testTable_REALTIME";
    String instanceId = "server01";
    Map<String, Map<String, String>> segmentAssignment = new HashMap<>();
    Map<String, SegmentZKMetadata> segmentMetadataMap = new HashMap<>();
    Set<String> preloadedSegments = new HashSet<>();
    AtomicBoolean wasPreloading = new AtomicBoolean(false);
    TableDataManager tableDataManager = mock(TableDataManager.class);
    UpsertContext upsertContext = mock(UpsertContext.class);
    when(upsertContext.isSnapshotEnabled()).thenReturn(true);
    when(upsertContext.isPreloadEnabled()).thenReturn(true);
    when(upsertContext.getTableDataManager()).thenReturn(tableDataManager);
    DummyPartitionUpsertMetadataManager upsertMetadataManager =
        new DummyPartitionUpsertMetadataManager(realtimeTableName, 0, upsertContext) {

          @Override
          Map<String, Map<String, String>> getSegmentAssignment(HelixManager helixManager) {
            return segmentAssignment;
          }

          @Override
          Map<String, SegmentZKMetadata> getSegmentsZKMetadata(HelixManager helixManager) {
            return segmentMetadataMap;
          }

          @Override
          void doPreloadSegmentWithSnapshot(TableDataManager tableDataManager, String segmentName,
              IndexLoadingConfig indexLoadingConfig, SegmentZKMetadata segmentZKMetadata) {
            wasPreloading.set(isPreloading());
            preloadedSegments.add(segmentName);
          }
        };

    // Setup mocks for TableConfig and Schema.
    TableConfig tableConfig = mock(TableConfig.class);
    UpsertConfig upsertConfig = new UpsertConfig();
    upsertConfig.setComparisonColumn("ts");
    upsertConfig.setEnablePreload(true);
    upsertConfig.setEnableSnapshot(true);
    when(tableConfig.getUpsertConfig()).thenReturn(upsertConfig);
    when(tableConfig.getTableName()).thenReturn(realtimeTableName);
    Schema schema = mock(Schema.class);
    when(schema.getPrimaryKeyColumns()).thenReturn(Collections.singletonList("pk"));
    IndexLoadingConfig indexLoadingConfig = mock(IndexLoadingConfig.class);
    when(indexLoadingConfig.getTableConfig()).thenReturn(tableConfig);

    // Setup mocks for HelixManager.
    HelixManager helixManager = mock(HelixManager.class);
    ZkHelixPropertyStore<ZNRecord> propertyStore = mock(ZkHelixPropertyStore.class);
    when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore);

    // Setup segment assignment. Only ONLINE segments are preloaded.
    segmentAssignment.put("consuming_seg01", ImmutableMap.of(instanceId, "CONSUMING"));
    segmentAssignment.put("consuming_seg02", ImmutableMap.of(instanceId, "CONSUMING"));
    segmentAssignment.put("offline_seg01", ImmutableMap.of(instanceId, "OFFLINE"));
    segmentAssignment.put("offline_seg02", ImmutableMap.of(instanceId, "OFFLINE"));
    String seg01Name = "testTable__0__1__" + System.currentTimeMillis();
    segmentAssignment.put(seg01Name, ImmutableMap.of(instanceId, "ONLINE"));
    String seg02Name = "testTable__0__2__" + System.currentTimeMillis();
    segmentAssignment.put(seg02Name, ImmutableMap.of(instanceId, "ONLINE"));
    // This segment is skipped as it's not from partition 0.
    String seg03Name = "testTable__1__3__" + System.currentTimeMillis();
    segmentAssignment.put(seg03Name, ImmutableMap.of(instanceId, "ONLINE"));

    SegmentZKMetadata zkMetadata = new SegmentZKMetadata(seg01Name);
    zkMetadata.setStatus(CommonConstants.Segment.Realtime.Status.DONE);
    segmentMetadataMap.put(seg01Name, zkMetadata);
    zkMetadata = new SegmentZKMetadata(seg02Name);
    zkMetadata.setStatus(CommonConstants.Segment.Realtime.Status.DONE);
    segmentMetadataMap.put(seg02Name, zkMetadata);
    zkMetadata = new SegmentZKMetadata(seg03Name);
    zkMetadata.setStatus(CommonConstants.Segment.Realtime.Status.DONE);
    segmentMetadataMap.put(seg03Name, zkMetadata);

    // Setup mocks to get file path to validDocIds snapshot.
    ExecutorService segmentPreloadExecutor = Executors.newFixedThreadPool(1);
    File tableDataDir = new File(TEMP_DIR, realtimeTableName);
    when(tableDataManager.getHelixManager()).thenReturn(helixManager);
    when(tableDataManager.getSegmentPreloadExecutor()).thenReturn(segmentPreloadExecutor);
    when(tableDataManager.getTableDataDir()).thenReturn(tableDataDir);
    InstanceDataManagerConfig instanceDataManagerConfig = mock(InstanceDataManagerConfig.class);
    when(instanceDataManagerConfig.getInstanceId()).thenReturn(instanceId);
    when(tableDataManager.getInstanceDataManagerConfig()).thenReturn(instanceDataManagerConfig);

    // No snapshot file for seg01, so it's skipped.
    File seg01IdxDir = new File(tableDataDir, seg01Name);
    FileUtils.forceMkdir(seg01IdxDir);
    when(tableDataManager.getSegmentDataDir(seg01Name, null, tableConfig)).thenReturn(seg01IdxDir);

    File seg02IdxDir = new File(tableDataDir, seg02Name);
    FileUtils.forceMkdir(seg02IdxDir);
    FileUtils.touch(new File(new File(seg02IdxDir, "v3"), V1Constants.VALID_DOC_IDS_SNAPSHOT_FILE_NAME));
    when(tableDataManager.getSegmentDataDir(seg02Name, null, tableConfig)).thenReturn(seg02IdxDir);

    try {
      // If preloading is enabled, the _isPreloading flag is true initially, until preloading is done.
      assertTrue(upsertMetadataManager.isPreloading());
      upsertMetadataManager.preloadSegments(indexLoadingConfig);
      assertEquals(preloadedSegments.size(), 1);
      assertTrue(preloadedSegments.contains(seg02Name));
      assertTrue(wasPreloading.get());
      assertFalse(upsertMetadataManager.isPreloading());
    } finally {
      segmentPreloadExecutor.shutdownNow();
    }
  }