org.apache.helix.store.zk.ZkHelixPropertyStore Java Examples

public static TableDataManager getTableDataManager(@Nonnull TableDataManagerConfig tableDataManagerConfig,
    @Nonnull String instanceId, @Nonnull ZkHelixPropertyStore<ZNRecord> propertyStore,
    @Nonnull ServerMetrics serverMetrics) {
  TableDataManager tableDataManager;
  switch (TableType.valueOf(tableDataManagerConfig.getTableDataManagerType())) {
    case OFFLINE:
      tableDataManager = new OfflineTableDataManager();
      break;
    case REALTIME:
      tableDataManager = new RealtimeTableDataManager(_segmentBuildSemaphore);
      break;
    default:
      throw new IllegalStateException();
  }
  tableDataManager.init(tableDataManagerConfig, instanceId, propertyStore, serverMetrics);
  return tableDataManager;
}
 

private TableDataManager makeTestableManager()
    throws Exception {
  TableDataManager tableDataManager = new OfflineTableDataManager();
  TableDataManagerConfig config;
  {
    config = mock(TableDataManagerConfig.class);
    when(config.getTableName()).thenReturn(TABLE_NAME);
    when(config.getDataDir()).thenReturn(_tmpDir.getAbsolutePath());
  }
  tableDataManager
      .init(config, "dummyInstance", mock(ZkHelixPropertyStore.class), new ServerMetrics(new MetricsRegistry()));
  tableDataManager.start();
  Field segsMapField = BaseTableDataManager.class.getDeclaredField("_segmentDataManagerMap");
  segsMapField.setAccessible(true);
  _internalSegMap = (Map<String, ImmutableSegmentDataManager>) segsMapField.get(tableDataManager);
  return tableDataManager;
}
 

/**
 * NOTE: this method is very expensive, use {@link #getSegments(ZkHelixPropertyStore, String)} instead if only segment
 * segment names are needed.
 */
public static List<OfflineSegmentZKMetadata> getOfflineSegmentZKMetadataListForTable(
    ZkHelixPropertyStore<ZNRecord> propertyStore, String tableName) {
  String offlineTableName = TableNameBuilder.OFFLINE.tableNameWithType(tableName);
  String parentPath = constructPropertyStorePathForResource(offlineTableName);
  List<ZNRecord> znRecords = propertyStore.getChildren(parentPath, null, AccessOption.PERSISTENT);
  if (znRecords != null) {
    int numZNRecords = znRecords.size();
    List<OfflineSegmentZKMetadata> offlineSegmentZKMetadataList = new ArrayList<>(numZNRecords);
    for (ZNRecord znRecord : znRecords) {
      // NOTE: it is possible that znRecord is null if the record gets removed while calling this method
      if (znRecord != null) {
        offlineSegmentZKMetadataList.add(new OfflineSegmentZKMetadata(znRecord));
      }
    }
    int numNullZNRecords = numZNRecords - offlineSegmentZKMetadataList.size();
    if (numNullZNRecords > 0) {
      LOGGER.warn("Failed to read {}/{} offline segment ZK metadata under path: {}", numZNRecords - numNullZNRecords,
          numZNRecords, parentPath);
    }
    return offlineSegmentZKMetadataList;
  } else {
    LOGGER.warn("Path: {} does not exist", parentPath);
    return Collections.emptyList();
  }
}
 

@Nullable
public static RealtimeSegmentZKMetadata getRealtimeSegmentZKMetadata(
    @Nonnull ZkHelixPropertyStore<ZNRecord> propertyStore, @Nonnull String tableName, @Nonnull String segmentName) {
  String realtimeTableName = TableNameBuilder.REALTIME.tableNameWithType(tableName);
  ZNRecord znRecord = propertyStore
      .get(constructPropertyStorePathForSegment(realtimeTableName, segmentName), null, AccessOption.PERSISTENT);
  // It is possible that the segment metadata has just been deleted due to retention.
  if (znRecord == null) {
    return null;
  }
  if (SegmentName.isHighLevelConsumerSegmentName(segmentName)) {
    return new RealtimeSegmentZKMetadata(znRecord);
  } else {
    return new LLCRealtimeSegmentZKMetadata(znRecord);
  }
}
 

/**
 * Instantiate a MockHelixManager.
 * @param instanceName the name of the instance associated with this manager.
 * @param instanceType the {@link InstanceType} of the requester.
 * @param zkAddr the address identifying the zk service to which this request is to be made.
 * @param helixCluster the {@link MockHelixCluster} associated with this manager.
 * @param znRecordMap a map that contains ZNode path and its {@link ZNRecord} associated with HelixPropertyStore in this manager.
 * @param beBadException the {@link Exception} that this manager will throw on listener registrations.
 */
MockHelixManager(String instanceName, InstanceType instanceType, String zkAddr, MockHelixCluster helixCluster,
    Map<String, ZNRecord> znRecordMap, Exception beBadException) {
  this.instanceName = instanceName;
  this.instanceType = instanceType;
  mockAdmin = helixCluster.getHelixAdminFactory().getHelixAdmin(zkAddr);
  mockAdmin.addHelixManager(this);
  clusterName = helixCluster.getClusterName();
  dataAccessor = new MockHelixDataAccessor(clusterName, mockAdmin);
  this.beBadException = beBadException;
  this.znRecordMap = znRecordMap;
  Properties storeProps = new Properties();
  storeProps.setProperty("helix.property.store.root.path",
      "/" + clusterName + "/" + ClusterMapUtils.PROPERTYSTORE_STR);
  HelixPropertyStoreConfig propertyStoreConfig = new HelixPropertyStoreConfig(new VerifiableProperties(storeProps));
  helixPropertyStore =
      (ZkHelixPropertyStore<ZNRecord>) CommonUtils.createHelixPropertyStore(zkAddr, propertyStoreConfig,
          Collections.singletonList(propertyStoreConfig.rootPath));
  if (znRecordMap != null) {
    for (Map.Entry<String, ZNRecord> znodePathAndRecord : znRecordMap.entrySet()) {
      helixPropertyStore.set(znodePathAndRecord.getKey(), znodePathAndRecord.getValue(), AccessOption.PERSISTENT);
    }
  }
}
 

@Override
public boolean execute()
    throws Exception {
  if (!_validateTableConfig && !_validateSchema) {
    throw new RuntimeException("Need to specify at least one of -schema and -tableConfig");
  }

  LOGGER.info("Connecting to Zookeeper: {}, cluster: ", _zkAddress, _clusterName);
  ZNRecordSerializer serializer = new ZNRecordSerializer();
  String path = PropertyPathConfig.getPath(PropertyType.PROPERTYSTORE, _clusterName);
  _helixPropertyStore = new ZkHelixPropertyStore<>(_zkAddress, serializer, path);

  LOGGER.info("\n\n-------------------- Starting Validation --------------------");
  if (_validateTableConfig) {
    validateTableConfig();
  }
  if (_validateSchema) {
    validateSchema();
  }

  return true;
}
 

/**
 * Fetches the instance partitions from Helix property store if it exists, or computes it for backward-compatibility.
 */
public static InstancePartitions fetchOrComputeInstancePartitions(HelixManager helixManager, TableConfig tableConfig,
    InstancePartitionsType instancePartitionsType) {
  String tableNameWithType = tableConfig.getTableName();

  // Fetch the instance partitions from property store if it exists
  ZkHelixPropertyStore<ZNRecord> propertyStore = helixManager.getHelixPropertyStore();
  InstancePartitions instancePartitions =
      fetchInstancePartitions(propertyStore, getInstancePartitionsName(tableNameWithType, instancePartitionsType));
  if (instancePartitions != null) {
    return instancePartitions;
  }

  // Compute the default instance partitions (for backward-compatibility)
  return computeDefaultInstancePartitions(helixManager, tableConfig, instancePartitionsType);
}
 

@Nullable
private static PartitionSegmentPruner getPartitionSegmentPruner(TableConfig tableConfig,
    ZkHelixPropertyStore<ZNRecord> propertyStore) {
  String tableNameWithType = tableConfig.getTableName();
  SegmentPartitionConfig segmentPartitionConfig = tableConfig.getIndexingConfig().getSegmentPartitionConfig();
  if (segmentPartitionConfig == null) {
    LOGGER.warn("Cannot enable partition pruning without segment partition config for table: {}", tableNameWithType);
    return null;
  }
  Map<String, ColumnPartitionConfig> columnPartitionMap = segmentPartitionConfig.getColumnPartitionMap();
  if (columnPartitionMap.size() != 1) {
    LOGGER.warn("Cannot enable partition pruning with other than exact one partition column for table: {}",
        tableNameWithType);
    return null;
  } else {
    String partitionColumn = columnPartitionMap.keySet().iterator().next();
    LOGGER.info("Using PartitionSegmentPruner on partition column: {} for table: {}", partitionColumn,
        tableNameWithType);
    return new PartitionSegmentPruner(tableNameWithType, partitionColumn, propertyStore);
  }
}
 

/**
 * NOTE: this method is very expensive, use {@link #getSegments(ZkHelixPropertyStore, String)} instead if only segment
 * segment names are needed.
 */
public static List<RealtimeSegmentZKMetadata> getRealtimeSegmentZKMetadataListForTable(
    ZkHelixPropertyStore<ZNRecord> propertyStore, String tableName) {
  String realtimeTableName = TableNameBuilder.REALTIME.tableNameWithType(tableName);
  String parentPath = constructPropertyStorePathForResource(realtimeTableName);
  List<ZNRecord> znRecords = propertyStore.getChildren(parentPath, null, AccessOption.PERSISTENT);
  if (znRecords != null) {
    int numZNRecords = znRecords.size();
    List<RealtimeSegmentZKMetadata> realtimeSegmentZKMetadataList = new ArrayList<>(numZNRecords);
    for (ZNRecord znRecord : znRecords) {
      // NOTE: it is possible that znRecord is null if the record gets removed while calling this method
      if (znRecord != null) {
        realtimeSegmentZKMetadataList.add(new RealtimeSegmentZKMetadata(znRecord));
      }
    }
    int numNullZNRecords = numZNRecords - realtimeSegmentZKMetadataList.size();
    if (numNullZNRecords > 0) {
      LOGGER.warn("Failed to read {}/{} realtime segment ZK metadata under path: {}", numZNRecords - numNullZNRecords,
          numZNRecords, parentPath);
    }
    return realtimeSegmentZKMetadataList;
  } else {
    LOGGER.warn("Path: {} does not exist", parentPath);
    return Collections.emptyList();
  }
}
 

/**
 * Returns the segments for the given table.
 *
 * @param propertyStore Helix property store
 * @param tableNameWithType Table name with type suffix
 * @return List of segment names
 */
public static List<String> getSegments(ZkHelixPropertyStore<ZNRecord> propertyStore, String tableNameWithType) {
  String segmentsPath = constructPropertyStorePathForResource(tableNameWithType);
  if (propertyStore.exists(segmentsPath, AccessOption.PERSISTENT)) {
    return propertyStore.getChildNames(segmentsPath, AccessOption.PERSISTENT);
  } else {
    return Collections.emptyList();
  }
}
 

@SuppressWarnings("unchecked")
protected void addTable(String tableNameWithType) {
  TableDataManagerConfig tableDataManagerConfig = mock(TableDataManagerConfig.class);
  when(tableDataManagerConfig.getTableDataManagerType())
      .thenReturn(TableNameBuilder.getTableTypeFromTableName(tableNameWithType).name());
  when(tableDataManagerConfig.getTableName()).thenReturn(tableNameWithType);
  when(tableDataManagerConfig.getDataDir()).thenReturn(INDEX_DIR.getAbsolutePath());
  TableDataManager tableDataManager = TableDataManagerProvider
      .getTableDataManager(tableDataManagerConfig, "testInstance", mock(ZkHelixPropertyStore.class),
          mock(ServerMetrics.class));
  tableDataManager.start();
  _tableDataManagerMap.put(tableNameWithType, tableDataManager);
}
 

/**
 * Read the segment lineage from the property store
 *
 * @param propertyStore  a property store
 * @param tableNameWithType a table name with type
 * @return a segment lineage, return null if znode does not exist
 */
public static SegmentLineage getSegmentLineage(ZkHelixPropertyStore<ZNRecord> propertyStore,
    String tableNameWithType) {
  ZNRecord znRecord = getSegmentLineageZNRecord(propertyStore, tableNameWithType);
  SegmentLineage segmentMergeLineage = null;
  if (znRecord != null) {
    segmentMergeLineage = SegmentLineage.fromZNRecord(znRecord);
  }
  return segmentMergeLineage;
}
 

/**
 * Read the segment lineage ZNRecord from the property store
 *
 * @param propertyStore a property store
 * @param tableNameWithType a table name with type
 * @return a ZNRecord of segment merge lineage, return null if znode does not exist
 */
public static ZNRecord getSegmentLineageZNRecord(ZkHelixPropertyStore<ZNRecord> propertyStore,
    String tableNameWithType) {
  String path = ZKMetadataProvider.constructPropertyStorePathForSegmentLineage(tableNameWithType);
  Stat stat = new Stat();
  ZNRecord segmentLineageZNRecord = propertyStore.get(path, stat, AccessOption.PERSISTENT);
  if (segmentLineageZNRecord != null) {
    segmentLineageZNRecord.setVersion(stat.getVersion());
  }
  return segmentLineageZNRecord;
}
 

@Nullable
public static Schema getSchema(@Nonnull ZkHelixPropertyStore<ZNRecord> propertyStore, @Nonnull String schemaName) {
  try {
    ZNRecord schemaZNRecord =
        propertyStore.get(constructPropertyStorePathForSchema(schemaName), null, AccessOption.PERSISTENT);
    if (schemaZNRecord == null) {
      return null;
    }
    return SchemaUtils.fromZNRecord(schemaZNRecord);
  } catch (Exception e) {
    LOGGER.error("Caught exception while getting schema: {}", schemaName, e);
    return null;
  }
}
 

public static void removeResourceSegmentsFromPropertyStore(ZkHelixPropertyStore<ZNRecord> propertyStore,
    String resourceName) {
  String propertyStorePath = constructPropertyStorePathForResource(resourceName);
  if (propertyStore.exists(propertyStorePath, AccessOption.PERSISTENT)) {
    propertyStore.remove(propertyStorePath, AccessOption.PERSISTENT);
  }
}
 

public static void removeResourceConfigFromPropertyStore(ZkHelixPropertyStore<ZNRecord> propertyStore,
    String resourceName) {
  String propertyStorePath = constructPropertyStorePathForResourceConfig(resourceName);
  if (propertyStore.exists(propertyStorePath, AccessOption.PERSISTENT)) {
    propertyStore.remove(propertyStorePath, AccessOption.PERSISTENT);
  }
}
 

public static boolean setOfflineSegmentZKMetadata(ZkHelixPropertyStore<ZNRecord> propertyStore,
    String offlineTableName, OfflineSegmentZKMetadata offlineSegmentZKMetadata, int expectedVersion) {
  // NOTE: Helix will throw ZkBadVersionException if version does not match
  try {
    return propertyStore
        .set(constructPropertyStorePathForSegment(offlineTableName, offlineSegmentZKMetadata.getSegmentName()),
            offlineSegmentZKMetadata.toZNRecord(), expectedVersion, AccessOption.PERSISTENT);
  } catch (ZkBadVersionException e) {
    return false;
  }
}
 

@Nullable
public static ZNRecord getZnRecord(@Nonnull ZkHelixPropertyStore<ZNRecord> propertyStore, @Nonnull String path) {
  Stat stat = new Stat();
  ZNRecord znRecord = propertyStore.get(path, stat, AccessOption.PERSISTENT);
  if (znRecord != null) {
    znRecord.setCreationTime(stat.getCtime());
    znRecord.setModifiedTime(stat.getMtime());
    znRecord.setVersion(stat.getVersion());
  }
  return znRecord;
}
 

@Nullable
public static OfflineSegmentZKMetadata getOfflineSegmentZKMetadata(
    @Nonnull ZkHelixPropertyStore<ZNRecord> propertyStore, @Nonnull String tableName, @Nonnull String segmentName) {
  String offlineTableName = TableNameBuilder.OFFLINE.tableNameWithType(tableName);
  ZNRecord znRecord = propertyStore
      .get(constructPropertyStorePathForSegment(offlineTableName, segmentName), null, AccessOption.PERSISTENT);
  if (znRecord == null) {
    return null;
  }
  return new OfflineSegmentZKMetadata(znRecord);
}
 

/**
 * Returns the LLC realtime segments for the given table.
 *
 * @param propertyStore Helix property store
 * @param realtimeTableName Realtime table name
 * @return List of LLC realtime segment names
 */
public static List<String> getLLCRealtimeSegments(ZkHelixPropertyStore<ZNRecord> propertyStore,
    String realtimeTableName) {
  List<String> llcRealtimeSegments = new ArrayList<>();
  String segmentsPath = constructPropertyStorePathForResource(realtimeTableName);
  if (propertyStore.exists(segmentsPath, AccessOption.PERSISTENT)) {
    List<String> segments = propertyStore.getChildNames(segmentsPath, AccessOption.PERSISTENT);
    for (String segment : segments) {
      if (SegmentName.isLowLevelConsumerSegmentName(segment)) {
        llcRealtimeSegments.add(segment);
      }
    }
  }
  return llcRealtimeSegments;
}
 

public BaseBrokerRequestHandler(Configuration config, RoutingManager routingManager,
    AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, BrokerMetrics brokerMetrics,
    ZkHelixPropertyStore<ZNRecord> propertyStore) {
  _config = config;
  _routingManager = routingManager;
  _accessControlFactory = accessControlFactory;
  _queryQuotaManager = queryQuotaManager;
  _brokerMetrics = brokerMetrics;

  _enableCaseInsensitive = _config.getBoolean(CommonConstants.Helix.ENABLE_CASE_INSENSITIVE_KEY, false);
  if (_enableCaseInsensitive) {
    _tableCache = new TableCache(propertyStore);
  } else {
    _tableCache = null;
  }
  _defaultHllLog2m = _config
      .getInt(CommonConstants.Helix.DEFAULT_HYPERLOGLOG_LOG2M_KEY, CommonConstants.Helix.DEFAULT_HYPERLOGLOG_LOG2M);

  _enableQueryLimitOverride = _config.getBoolean(Broker.CONFIG_OF_ENABLE_QUERY_LIMIT_OVERRIDE, false);

  _brokerId = config.getString(Broker.CONFIG_OF_BROKER_ID, getDefaultBrokerId());
  _brokerTimeoutMs = config.getLong(Broker.CONFIG_OF_BROKER_TIMEOUT_MS, Broker.DEFAULT_BROKER_TIMEOUT_MS);
  _queryResponseLimit =
      config.getInt(Broker.CONFIG_OF_BROKER_QUERY_RESPONSE_LIMIT, Broker.DEFAULT_BROKER_QUERY_RESPONSE_LIMIT);
  _queryLogLength = config.getInt(Broker.CONFIG_OF_BROKER_QUERY_LOG_LENGTH, Broker.DEFAULT_BROKER_QUERY_LOG_LENGTH);
  _queryLogRateLimiter = RateLimiter.create(config.getDouble(Broker.CONFIG_OF_BROKER_QUERY_LOG_MAX_RATE_PER_SECOND,
      Broker.DEFAULT_BROKER_QUERY_LOG_MAX_RATE_PER_SECOND));
  _numDroppedLog = new AtomicInteger(0);
  _numDroppedLogRateLimiter = RateLimiter.create(1.0);

  LOGGER
      .info("Broker Id: {}, timeout: {}ms, query response limit: {}, query log length: {}, query log max rate: {}qps",
          _brokerId, _brokerTimeoutMs, _queryResponseLimit, _queryLogLength, _queryLogRateLimiter.getRate());
}
 

/**
 * Get the schema associated with the given table name.
 *
 * @param propertyStore Helix property store
 * @param tableName Table name with or without type suffix.
 * @return Schema associated with the given table name.
 */
@Nullable
public static Schema getTableSchema(@Nonnull ZkHelixPropertyStore<ZNRecord> propertyStore,
    @Nonnull String tableName) {
  String rawTableName = TableNameBuilder.extractRawTableName(tableName);
  Schema schema = getSchema(propertyStore, rawTableName);
  if (schema != null) {
    return schema;
  }

  // For backward compatible where schema name is not the same as raw table name
  TableType tableType = TableNameBuilder.getTableTypeFromTableName(tableName);
  // Try to fetch realtime schema first
  if (tableType == null || tableType == TableType.REALTIME) {
    TableConfig realtimeTableConfig = getRealtimeTableConfig(propertyStore, tableName);
    if (realtimeTableConfig != null) {
      schema = getSchema(propertyStore, realtimeTableConfig.getValidationConfig().getSchemaName());
    }
  }
  // Try to fetch offline schema if realtime schema does not exist
  if (schema == null && (tableType == null || tableType == TableType.OFFLINE)) {
    schema = getSchema(propertyStore, TableNameBuilder.OFFLINE.tableNameWithType(tableName));
  }
  if (schema != null) {
    LOGGER.warn("Schema name does not match raw table name, schema name: {}, raw table name: {}",
        schema.getSchemaName(), TableNameBuilder.extractRawTableName(tableName));
  }
  return schema;
}
 

public PartitionSegmentPruner(String tableNameWithType, String partitionColumn,
    ZkHelixPropertyStore<ZNRecord> propertyStore) {
  _tableNameWithType = tableNameWithType;
  _partitionColumn = partitionColumn;
  _propertyStore = propertyStore;
  _segmentZKMetadataPathPrefix = ZKMetadataProvider.constructPropertyStorePathForResource(tableNameWithType) + "/";
}
 

public TimeBoundaryManager(TableConfig tableConfig, ZkHelixPropertyStore<ZNRecord> propertyStore) {
  Preconditions.checkState(tableConfig.getTableType() == TableType.OFFLINE,
      "Cannot construct TimeBoundaryManager for real-time table: %s", tableConfig.getTableName());
  _offlineTableName = tableConfig.getTableName();
  _propertyStore = propertyStore;
  _segmentZKMetadataPathPrefix = ZKMetadataProvider.constructPropertyStorePathForResource(_offlineTableName) + "/";

  Schema schema = ZKMetadataProvider.getTableSchema(_propertyStore, _offlineTableName);
  Preconditions.checkState(schema != null, "Failed to find schema for table: %s", _offlineTableName);
  _timeColumn = tableConfig.getValidationConfig().getTimeColumnName();
  Preconditions
      .checkNotNull(_timeColumn, "Time column must be configured in table config for table: %s", _offlineTableName);
  DateTimeFieldSpec dateTimeSpec = schema.getSpecForTimeColumn(_timeColumn);
  Preconditions.checkNotNull(dateTimeSpec, "Field spec must be specified in schema for time column: %s of table: %s",
      _timeColumn, _offlineTableName);
  DateTimeFormatSpec formatSpec = new DateTimeFormatSpec(dateTimeSpec.getFormat());
  _timeUnit = formatSpec.getColumnUnit();
  Preconditions
      .checkNotNull(_timeUnit, "Time unit must be configured in the field spec for time column: %s of table: %s",
          _timeColumn, _offlineTableName);

  // For HOURLY table with time unit other than DAYS, use (maxEndTime - 1 HOUR) as the time boundary; otherwise, use
  // (maxEndTime - 1 DAY)
  _isHourlyTable = "HOURLY".equalsIgnoreCase(tableConfig.getValidationConfig().getSegmentPushFrequency())
      && _timeUnit != TimeUnit.DAYS;

  LOGGER.info("Constructed TimeBoundaryManager with timeColumn: {}, timeUnit: {}, isHourlyTable: {} for table: {}",
      _timeColumn, _timeUnit, _isHourlyTable, _offlineTableName);
}
 

private void testAllFailed(List<String> segments) {
  HelixAdmin helixAdmin = makeHelixAdmin();
  ZkHelixPropertyStore<ZNRecord> propertyStore = makePropertyStore();
  FakeDeletionManager deletionManager = new FakeDeletionManager(helixAdmin, propertyStore);
  deletionManager.deleteSegmentsFromPropertyStoreAndLocal(tableName, segments);

  Assert.assertTrue(deletionManager.segmentsToRetry.containsAll(segments));
  Assert.assertEquals(deletionManager.segmentsToRetry.size(), segments.size());
  Assert.assertEquals(deletionManager.segmentsRemovedFromStore.size(), 0);
}
 

@Test
public void allPassed() {
  HelixAdmin helixAdmin = makeHelixAdmin();
  ZkHelixPropertyStore<ZNRecord> propertyStore = makePropertyStore();
  FakeDeletionManager deletionManager = new FakeDeletionManager(helixAdmin, propertyStore);
  Set<String> segments = new HashSet<>();
  segments.addAll(segmentsThatShouldBeDeleted());
  deletionManager.deleteSegmentsFromPropertyStoreAndLocal(tableName, segments);

  Assert.assertEquals(deletionManager.segmentsToRetry.size(), 0);
  Assert.assertEquals(deletionManager.segmentsRemovedFromStore.size(), segments.size());
  Assert.assertTrue(deletionManager.segmentsRemovedFromStore.containsAll(segments));
}
 

public static boolean getClusterTenantIsolationEnabled(ZkHelixPropertyStore<ZNRecord> propertyStore) {
  String controllerConfigPath = constructPropertyStorePathForControllerConfig(CLUSTER_TENANT_ISOLATION_ENABLED_KEY);
  if (propertyStore.exists(controllerConfigPath, AccessOption.PERSISTENT)) {
    ZNRecord znRecord = propertyStore.get(controllerConfigPath, null, AccessOption.PERSISTENT);
    if (znRecord.getSimpleFields().containsKey(CLUSTER_TENANT_ISOLATION_ENABLED_KEY)) {
      return znRecord.getBooleanField(CLUSTER_TENANT_ISOLATION_ENABLED_KEY, true);
    } else {
      return true;
    }
  } else {
    return true;
  }
}
 

ZkHelixPropertyStore<ZNRecord> makePropertyStore() {
  ZkHelixPropertyStore store = mock(ZkHelixPropertyStore.class);
  final List<String> failedSegs = segmentsFailingPropStore();
  /*
  for (String segment : failedSegs) {
    String propStorePath = ZKMetadataProvider.constructPropertyStorePathForSegment(tableName, segment);
    when(store.remove(propStorePath, AccessOption.PERSISTENT)).thenReturn(false);
  }
  List<String> successfulSegs = segmentsThatShouldBeDeleted();
  for (String segment : successfulSegs) {
    String propStorePath = ZKMetadataProvider.constructPropertyStorePathForSegment(tableName, segment);
    when(store.remove(propStorePath, AccessOption.PERSISTENT)).thenReturn(true);
  }
  */
  when(store.remove(anyList(), anyInt())).thenAnswer(new Answer<boolean[]>() {
    @Override
    public boolean[] answer(InvocationOnMock invocationOnMock)
        throws Throwable {
      List<String> propStoreList = (List<String>) (invocationOnMock.getArguments()[0]);
      boolean[] result = new boolean[propStoreList.size()];
      for (int i = 0; i < result.length; i++) {
        final String path = propStoreList.get(i);
        final String segmentId = path.substring((path.lastIndexOf('/') + 1));
        if (failedSegs.indexOf(segmentId) < 0) {
          result[i] = true;
        } else {
          result[i] = false;
        }
      }
      return result;
    }
  });

  when(store.exists(anyString(), anyInt())).thenReturn(true);
  return store;
}
 

public AutoAddInvertedIndex(@Nonnull String zkAddress, @Nonnull String clusterName, @Nonnull String controllerAddress,
    @Nonnull String brokerAddress, @Nonnull Strategy strategy, @Nonnull Mode mode) {
  _clusterName = clusterName;
  _controllerAddress = controllerAddress;
  _brokerAddress = brokerAddress;
  _helixAdmin = new ZKHelixAdmin(zkAddress);
  _propertyStore = new ZkHelixPropertyStore<>(zkAddress, new ZNRecordSerializer(),
      PropertyPathConfig.getPath(PropertyType.PROPERTYSTORE, clusterName));
  _strategy = strategy;
  _mode = mode;
}
 

@BeforeClass
public void setUp() {
  _zkInstance = ZkStarter.startLocalZkServer();
  _zkClient =
      new ZkClient(ZkStarter.DEFAULT_ZK_STR, ZkClient.DEFAULT_SESSION_TIMEOUT, ZkClient.DEFAULT_CONNECTION_TIMEOUT,
          new ZNRecordSerializer());
  _propertyStore =
      new ZkHelixPropertyStore<>(new ZkBaseDataAccessor<>(_zkClient), "/SegmentPrunerTest/PROPERTYSTORE", null);
}