org.apache.helix.AccessOption Java Examples

/**
 * 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();
  }
}
 

@Override
public boolean update(String path, DataUpdater<T> updater, int options) {
  if (_fallbackStore == null) {
    return super.update(path, updater, options);
  } else {
    Stat stat = super.getStat(path, options);
    if (stat == null) {
      // create znode at new location with fallback-value
      T fallbackValue = _fallbackStore.get(path, null, options);
      boolean succeed = super.create(path, fallbackValue, AccessOption.PERSISTENT);
      if (!succeed) {
        LOG.error("Can't update " + path + " since there are concurrent updates");
        return false;
      }
    }
    return super.update(path, updater, options);
  }
}
 

@Test(dependsOnMethods = "testCustomizedViewWithoutType")
public void testCustomizedViewCorrectConstructor() throws Exception {
  Map<PropertyType, List<String>> sourceDataTypes = new HashMap<>();
  sourceDataTypes.put(PropertyType.CUSTOMIZEDVIEW, Arrays.asList("typeA"));
  MockRoutingTableProvider routingTableProvider =
      new MockRoutingTableProvider(_spectator, sourceDataTypes);

  CustomizedView customizedView = new CustomizedView(TEST_DB);
  customizedView.setState("p1", "h1", "testState");

  // Clear the flag before writing to the Customized View Path
  customizedViewChangeCalled.getAndSet(false);
  String customizedViewPath = PropertyPathBuilder.customizedView(CLUSTER_NAME, "typeA", TEST_DB);
  _spectator.getHelixDataAccessor().getBaseDataAccessor().set(customizedViewPath,
      customizedView.getRecord(), AccessOption.PERSISTENT);

  boolean onCustomizedViewChangeCalled =
      TestHelper.verify(() -> customizedViewChangeCalled.get(), WAIT_DURATION);
  Assert.assertTrue(onCustomizedViewChangeCalled);

  _spectator.getHelixDataAccessor().getBaseDataAccessor().remove(customizedViewPath,
      AccessOption.PERSISTENT);
  routingTableProvider.shutdown();
}
 

/**
 * This test method tests whether PropertyKey.Builder successfully creates a path for
 * WorkflowContext instances.
 * TODO: KeyBuilder must handle the case for future versions of Task Framework with a different
 * path structure
 */
@Test
public void testGetWorkflowContext() {
  // Manually create a WorkflowContext instance
  ZNRecord znRecord = new ZNRecord(WORKFLOW_NAME);
  WorkflowContext workflowContext = new WorkflowContext(znRecord);
  _manager.getHelixPropertyStore().set(
      Joiner.on("/").join(TaskConstants.REBALANCER_CONTEXT_ROOT, WORKFLOW_NAME, CONTEXT_NODE),
      workflowContext.getRecord(), AccessOption.PERSISTENT);

  // Test retrieving this WorkflowContext using PropertyKey.Builder.getPath()
  String path = KEY_BUILDER.workflowContext(WORKFLOW_NAME).getPath();
  WorkflowContext workflowCtx =
      new WorkflowContext(_baseAccessor.get(path, null, AccessOption.PERSISTENT));

  Assert.assertEquals(workflowContext, workflowCtx);
}
 

private void updateMessageState(List<Message> readMsgs, HelixDataAccessor accessor,
    String instanceName) {
  Builder keyBuilder = accessor.keyBuilder();
  List<String> readMsgPaths = new ArrayList<>();
  List<DataUpdater<ZNRecord>> updaters = new ArrayList<>();
  for (Message msg : readMsgs) {
    readMsgPaths.add(msg.getKey(keyBuilder, instanceName).getPath());
    _knownMessageIds.add(msg.getId());
    /**
     * We use the updater to avoid race condition between writing message to zk as READ state and removing message after ST is done
     * If there is no message at this path, meaning the message is removed so we do not write the message
     */
    updaters.add(new DataUpdater<ZNRecord>() {
      @Override
      public ZNRecord update(ZNRecord currentData) {
        if (currentData == null) {
          LOG.warn("Message {} targets at {} has already been removed before it is set as READ on instance {}", msg.getId(), msg.getTgtName(), instanceName);
          return null;
        }
        return msg.getRecord();
      }
    });
  }
  accessor.updateChildren(readMsgPaths, updaters, AccessOption.PERSISTENT);
}
 

@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);
  }
}
 

@Override
public boolean createStateModelDef(StateModelDefinition stateModelDef) {
  String path = PropertyPathBuilder.stateModelDef(_clusterName, stateModelDef.getId());
  HelixProperty property =
      getProperty(new PropertyKey.Builder(_clusterName).stateModelDef(stateModelDef.getId()));

  // Set new StateModelDefinition if it is different from old one.
  if (property != null) {
    // StateModelDefinition need to be updated
    if (!new StateModelDefinition(property.getRecord()).equals(stateModelDef)) {
      return stateModelDef.isValid() && _baseDataAccessor
          .set(path, stateModelDef.getRecord(), AccessOption.PERSISTENT);
    }
  } else {
    // StateModeDefinition does not exist
    return stateModelDef.isValid() && _baseDataAccessor
        .create(path, stateModelDef.getRecord(), AccessOption.PERSISTENT);
  }
  // StateModelDefinition exists but not need to be updated
  return true;
}
 

/**
 * Uploads cluster config infos onto Helix PropertyStore.
 * @param adminInfosByDc the cluster admin information (overridden partitions, added replicas) grouped by DC that would
 *                       be applied to cluster.
 * @param clusterAdminType the type of cluster admin that would be uploaded (i.e. PartitionOverride, ReplicaAddition)
 * @param adminConfigZNodePath ZNode path of admin config associated with clusterAdminType.
 */
private void uploadClusterAdminInfos(Map<String, Map<String, Map<String, String>>> adminInfosByDc,
    String clusterAdminType, String adminConfigZNodePath) {
  Properties storeProps = new Properties();
  storeProps.setProperty("helix.property.store.root.path", "/" + clusterName + "/" + PROPERTYSTORE_STR);
  HelixPropertyStoreConfig propertyStoreConfig = new HelixPropertyStoreConfig(new VerifiableProperties(storeProps));
  for (Map.Entry<String, ClusterMapUtils.DcZkInfo> entry : dataCenterToZkAddress.entrySet()) {
    info("Uploading {} infos for datacenter {}.", clusterAdminType, entry.getKey());
    List<String> zkConnectStrs = entry.getValue().getZkConnectStrs();
    // The number of zk endpoints has been validated in the ctor of HelixBootstrapUpgradeUtil, no need to check it again
    HelixPropertyStore<ZNRecord> helixPropertyStore =
        CommonUtils.createHelixPropertyStore(zkConnectStrs.get(0), propertyStoreConfig, null);
    ZNRecord znRecord = new ZNRecord(clusterAdminType);
    znRecord.setMapFields(adminInfosByDc.get(entry.getKey()));
    if (!helixPropertyStore.set(adminConfigZNodePath, znRecord, AccessOption.PERSISTENT)) {
      logger.error("Failed to upload {} infos for datacenter {}", clusterAdminType, entry.getKey());
    }
  }
}
 

@Test
public void testSyncSet() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String testName = className + "_" + methodName;

  System.out.println("START " + testName + " at " + new Date(System.currentTimeMillis()));

  String path = String.format("/%s/%s", _rootPath, "msg_0");
  ZNRecord record = new ZNRecord("msg_0");
  BaseDataAccessor<ZNRecord> accessor = new ZkBaseDataAccessor<ZNRecord>(_gZkClient);

  boolean success = accessor.set(path, record, AccessOption.PERSISTENT);
  Assert.assertTrue(success);
  ZNRecord getRecord = _gZkClient.readData(path);
  Assert.assertNotNull(getRecord);
  Assert.assertEquals(getRecord.getId(), "msg_0");

  System.out.println("END " + testName + " at " + new Date(System.currentTimeMillis()));
}
 

private void batchUpdateData(HelixDataAccessor accessor, List<String> dataUpdateNames,
    Map<String, ZNRecord> dataMap, Set<String> dataToUpdate, TaskDataType taskDataType) {
  List<String> contextUpdatePaths = new ArrayList<>();
  List<ZNRecord> updatedData = new ArrayList<>();
  for (String resourceName : dataUpdateNames) {
    if (dataMap.get(resourceName) != null) {
      contextUpdatePaths.add(getTaskDataPath(resourceName, taskDataType));
      updatedData.add(dataMap.get(resourceName));
    }
  }

  boolean[] updateSuccess = accessor.getBaseDataAccessor()
      .setChildren(contextUpdatePaths, updatedData, AccessOption.PERSISTENT);

  for (int i = 0; i < updateSuccess.length; i++) {
    if (updateSuccess[i]) {
      dataToUpdate.remove(dataUpdateNames.get(i));
    } else {
      LogUtil.logWarn(LOG, _controlContextProvider.getClusterEventId(), String
          .format("Failed to update the %s for %s", taskDataType.name(), dataUpdateNames.get(i)));
    }
  }
}
 

@Test
public void testDefaultAccessorCreateCustomData() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String testName = className + "_" + methodName;
  System.out.println("START " + testName + " at " + new Date(System.currentTimeMillis()));

  String path = String.format("/%s/%s", _rootPath, "msg_0");

  ZkBaseDataAccessor defaultAccessor = new ZkBaseDataAccessor(ZK_ADDR);

  List<Integer> l0 = ImmutableList.of(1, 2, 3);
  boolean createResult = defaultAccessor.create(path, l0, AccessOption.PERSISTENT);
  // The result is expected to be false because the list is not ZNRecord
  Assert.assertFalse(createResult);
  createResult = defaultAccessor.create(path, new ZNRecord("test"), AccessOption.PERSISTENT);
  // The result is expected to be true
  Assert.assertTrue(createResult);

  defaultAccessor.close();
  System.out.println("END " + testName + " at " + new Date(System.currentTimeMillis()));
}
 

/**
 * {@inheritDoc}
 *
 * Returns {@code true} does not guarantee all the {@link TopicListener} will receive the message. It just indicates
 * the message has been successfully sent out.
 */
@Override
public boolean publish(String topic, String message) {
  if (topic == null) {
    throw new IllegalArgumentException("topic cannot be null");
  }
  if (message == null) {
    throw new IllegalArgumentException("message cannot be null");
  }
  String topicPath = getTopicPath(topic);
  ZNRecord record = new ZNRecord(topicPath);
  record.setSimpleField(MESSAGE_KEY, message);
  boolean res = helixStore.set(topicPath, record, AccessOption.PERSISTENT);
  if (res) {
    logger.trace("message={} has been published for topic={}", message, topic);
  } else {
    logger.error("failed to publish message={} for topic={}", message, topic);
  }
  return res;
}
 

/**
 * Fetch the list of {@link RouterStore.BlobIDAndVersion} from the helixStore.
 * @param count The expected number of elements in the list.
 * @return The list of {@link RouterStore.BlobIDAndVersion}.
 */
private List<RouterStore.BlobIDAndVersion> getBlobIDAndVersionInHelix(int count) {
  // Verify that ZNRecord contains the right data.
  ZNRecord record = helixStore.get(RouterStore.ACCOUNT_METADATA_BLOB_IDS_PATH, null, AccessOption.PERSISTENT);
  assertNotNull("ZNRecord missing after update", record);
  List<String> accountBlobs = record.getListField(RouterStore.ACCOUNT_METADATA_BLOB_IDS_LIST_KEY);
  assertNotNull("Blob ids are missing from ZNRecord", accountBlobs);
  // version also equals to the number of blobs
  assertEquals("Number of blobs mismatch", count, accountBlobs.size());

  List<RouterStore.BlobIDAndVersion> blobIDAndVersions = new ArrayList<>(count);
  for (String json : accountBlobs) {
    blobIDAndVersions.add(RouterStore.BlobIDAndVersion.fromJson(json));
  }
  return blobIDAndVersions;
}
 

/**
 * 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();
  }
}
 

@Test
public void testSyncExist() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String testName = className + "_" + methodName;

  System.out.println("START " + testName + " at " + new Date(System.currentTimeMillis()));

  String path = String.format("/%s/%s", _rootPath, "msg_0");
  ZNRecord record = new ZNRecord("msg_0");
  ZkBaseDataAccessor<ZNRecord> accessor = new ZkBaseDataAccessor<ZNRecord>(_gZkClient);

  boolean success = accessor.exists(path, 0);
  Assert.assertFalse(success);

  success = accessor.create(path, record, AccessOption.EPHEMERAL);
  Assert.assertTrue(success);

  success = accessor.exists(path, 0);
  Assert.assertTrue(success);

  System.out.println("END " + testName + " at " + new Date(System.currentTimeMillis()));

}
 

@Test
public void testSyncGetStat() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String testName = className + "_" + methodName;

  System.out.println("START " + testName + " at " + new Date(System.currentTimeMillis()));

  String path = String.format("/%s/%s", _rootPath, "msg_0");
  ZNRecord record = new ZNRecord("msg_0");
  ZkBaseDataAccessor<ZNRecord> accessor = new ZkBaseDataAccessor<ZNRecord>(_gZkClient);

  Stat stat = accessor.getStat(path, 0);
  Assert.assertNull(stat);

  boolean success = accessor.create(path, record, AccessOption.EPHEMERAL);
  Assert.assertTrue(success);

  stat = accessor.getStat(path, 0);
  Assert.assertNotNull(stat);
  Assert.assertEquals(stat.getVersion(), 0);
  Assert.assertNotSame(stat.getEphemeralOwner(), 0);

  System.out.println("END " + testName + " at " + new Date(System.currentTimeMillis()));

}
 

@Test
public void testCustomAccessorCreateZnRecord() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String testName = className + "_" + methodName;
  System.out.println("START " + testName + " at " + new Date(System.currentTimeMillis()));

  String path = String.format("/%s/%s", _rootPath, "msg_0");

  ZkBaseDataAccessor customDataAccessor = new ZkBaseDataAccessor(ZK_ADDR, LIST_SERIALIZER);
  boolean createResult = customDataAccessor.create(path, new ZNRecord("test"), AccessOption.PERSISTENT);
  // The result is expected to be false because the ZnRecord is not List
  Assert.assertFalse(createResult);
  createResult = customDataAccessor.create(path, ImmutableList.of(1, 2, 3), AccessOption.PERSISTENT);
  // The result is expected to be true
  Assert.assertTrue(createResult);

  customDataAccessor.close();
  System.out.println("END " + testName + " at " + new Date(System.currentTimeMillis()));
}
 

@Test
public void testZkClientMonitor() throws JMException {
  final String TEST_ROOT = "/test_root";

  ZkHelixPropertyStore<ZNRecord> store =
      new ZkHelixPropertyStore<>(ZK_ADDR, new SerializableSerializer(), TEST_ROOT);

  ObjectName name = MBeanRegistrar.buildObjectName(MonitorDomainNames.HelixZkClient.name(),
      ZkClientMonitor.MONITOR_TYPE, ZkHelixPropertyStore.MONITOR_TYPE,
      ZkClientMonitor.MONITOR_KEY, TEST_ROOT, ZkClientPathMonitor.MONITOR_PATH, "Root");
  MBeanServer beanServer = ManagementFactory.getPlatformMBeanServer();
  Assert.assertTrue(beanServer.isRegistered(name));

  store.getStat("/", AccessOption.PERSISTENT);

  Assert.assertEquals((long) beanServer.getAttribute(name, "ReadCounter"), 1);
}
 

static void addOrUpdateTaskUserContentMap(final HelixPropertyStore<ZNRecord> propertyStore,
    final String job, final String task, final Map<String, String> contentToAddOrUpdate) {
  if (job == null || task == null) {
    throw new IllegalArgumentException(
        "job and task must be not null when adding task user content");
  }
  String path =
      Joiner.on("/").join(TaskConstants.REBALANCER_CONTEXT_ROOT, job, USER_CONTENT_NODE);

  if (!propertyStore.update(path, new DataUpdater<ZNRecord>() {
    @Override
    public ZNRecord update(ZNRecord znRecord) {
      if (znRecord == null) {
        // This indicates that somehow the UserContentStore ZNode is missing
        // This should not happen, but if it is missing, create one
        znRecord = new ZNRecord(new ZNRecord(TaskUtil.USER_CONTENT_NODE));
      }
      if (znRecord.getMapField(task) == null) {
        znRecord.setMapField(task, new HashMap<String, String>());
      }
      znRecord.getMapField(task).putAll(contentToAddOrUpdate);
      return znRecord;
    }
  }, AccessOption.PERSISTENT)) {
    LOG.error("Failed to update the task UserContentStore for task {} in job {}", task, job);
  }
}
 

@Test
public void testExpiry() throws Exception {
  String jobName = "Expiry";
  long expiry = 1000;
  Map<String, String> commandConfig = ImmutableMap.of(MockTask.JOB_DELAY, String.valueOf(100));
  JobConfig.Builder jobBuilder = JobConfig.Builder.fromMap(WorkflowGenerator.DEFAULT_JOB_CONFIG);
  jobBuilder.setJobCommandConfigMap(commandConfig);

  Workflow flow = WorkflowGenerator.generateSingleJobWorkflowBuilder(jobName, jobBuilder)
      .setExpiry(expiry).build();

  _driver.start(flow);
  _driver.pollForWorkflowState(jobName, TaskState.IN_PROGRESS);

  // Running workflow should have config and context viewable through accessor
  HelixDataAccessor accessor = _manager.getHelixDataAccessor();
  PropertyKey workflowCfgKey = accessor.keyBuilder().resourceConfig(jobName);
  String workflowPropStoreKey =
      Joiner.on("/").join(TaskConstants.REBALANCER_CONTEXT_ROOT, jobName);

  // Ensure context and config exist
  Assert.assertTrue(
      _manager.getHelixPropertyStore().exists(workflowPropStoreKey, AccessOption.PERSISTENT));
  Assert.assertNotSame(accessor.getProperty(workflowCfgKey), null);

  // Wait for job to finish and expire
  _driver.pollForWorkflowState(jobName, TaskState.COMPLETED);
  Thread.sleep(expiry + 100);

  // Ensure workflow config and context were cleaned up by now
  Assert.assertFalse(
      _manager.getHelixPropertyStore().exists(workflowPropStoreKey, AccessOption.PERSISTENT));
  Assert.assertNull(accessor.getProperty(workflowCfgKey));
}
 

@Test
public void testSyncCreateWithCustomSerializer() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String testName = className + "_" + methodName;
  System.out.println("START " + testName + " at " + new Date(System.currentTimeMillis()));

  String path = String.format("/%s/%s", _rootPath, "msg_0");

  ZkBaseDataAccessor<List<Integer>> accessor = new ZkBaseDataAccessor<>(ZK_ADDR, LIST_SERIALIZER);

  List<Integer> l0 = ImmutableList.of(1, 2, 3);
  List<Integer> l1 = ImmutableList.of(4, 5, 6);
  boolean createResult = accessor.create(path, l0, AccessOption.PERSISTENT);
  Assert.assertTrue(createResult);

  List<Integer> data = (List<Integer>) accessor.get(path, null, AccessOption.PERSISTENT);
  Assert.assertEquals(data, l0);
  boolean setResult = accessor.set(path, l1, 0, AccessOption.PERSISTENT);
  Assert.assertTrue(setResult);

  data = (List<Integer>) accessor.get(path, null, AccessOption.PERSISTENT);
  Assert.assertEquals(data, l1);

  accessor.close();
  System.out.println("END " + testName + " at " + new Date(System.currentTimeMillis()));
}
 

@Test
public void testSyncCreate() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String testName = className + "_" + methodName;

  System.out.println("START " + testName + " at " + new Date(System.currentTimeMillis()));

  String path = String.format("/%s/%s", _rootPath, "msg_0");
  ZNRecord record = new ZNRecord("msg_0");
  ZkBaseDataAccessor<ZNRecord> accessor = new ZkBaseDataAccessor<>(_gZkClient);

  boolean success = accessor.create(path, record, AccessOption.PERSISTENT);
  Assert.assertTrue(success);
  ZNRecord getRecord = _gZkClient.readData(path);
  Assert.assertNotNull(getRecord);
  Assert.assertEquals(getRecord.getId(), "msg_0");

  record.setSimpleField("key0", "value0");
  success = accessor.create(path, record, AccessOption.PERSISTENT);
  Assert.assertFalse(success, "Should fail since node already exists");
  getRecord = _gZkClient.readData(path);
  Assert.assertNotNull(getRecord);
  Assert.assertEquals(getRecord.getSimpleFields().size(), 0);

  System.out.println("END " + testName + " at " + new Date(System.currentTimeMillis()));
}
 

@Nullable
public static TableConfig getTableConfig(@Nonnull ZkHelixPropertyStore<ZNRecord> propertyStore,
    @Nonnull String tableNameWithType) {
  ZNRecord znRecord = propertyStore
      .get(constructPropertyStorePathForResourceConfig(tableNameWithType), null, AccessOption.PERSISTENT);
  if (znRecord == null) {
    return null;
  }
  try {
    return TableConfigUtils.fromZNRecord(znRecord);
  } catch (Exception e) {
    LOGGER.error("Caught exception while getting table configuration for table: {}", tableNameWithType, e);
    return null;
  }
}
 

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;
  }
}
 

private void setLeader(HelixManager manager) throws Exception {
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  final LiveInstance leader = new LiveInstance(manager.getInstanceName());
  leader.setLiveInstance(ManagementFactory.getRuntimeMXBean().getName());
  leader.setSessionId(manager.getSessionId());
  leader.setHelixVersion(manager.getVersion());

  // Delete the current controller leader node so it will trigger leader election
  while (!manager.isLeader()) {
    accessor.getBaseDataAccessor()
        .remove(PropertyPathBuilder.controllerLeader(manager.getClusterName()),
            AccessOption.EPHEMERAL);
    Thread.sleep(50);
  }
}
 

/**
 * update workflow's property to remove jobs from JOB_STATES if there are already started.
 */
protected static boolean removeJobsState(final HelixPropertyStore<ZNRecord> propertyStore,
    final String workflow, final Set<String> jobs) {
  String contextPath =
      Joiner.on("/").join(TaskConstants.REBALANCER_CONTEXT_ROOT, workflow, TaskUtil.CONTEXT_NODE);

  // If the queue is not started, there is no JobState need to be removed.
  if (!propertyStore.exists(contextPath, 0)) {
    return true;
  }

  DataUpdater<ZNRecord> updater = new DataUpdater<ZNRecord>() {
    @Override
    public ZNRecord update(ZNRecord currentData) {
      if (currentData != null) {
        WorkflowContext workflowContext = new WorkflowContext(currentData);
        workflowContext.removeJobStates(jobs);
        workflowContext.removeJobStartTime(jobs);
        currentData = workflowContext.getRecord();
      }
      return currentData;
    }
  };
  if (!propertyStore.update(contextPath, updater, AccessOption.PERSISTENT)) {
    LOG.warn("Fail to remove job state for jobs " + jobs + " from workflow " + workflow);
    return false;
  }
  return true;
}
 

/**
 * Read aggregator checkpoint from zookeeper
 *
 * @return timestamp
 */
public long readCheckpoint(AggregationTaskRunner.AGGREGATOR_NAME aggregatorName) {
  String path = getCheckpointZKPath(aggregatorName);
  LOG.debug("Reading checkpoint at " + path);
  Stat stat = new Stat();
  ZNRecord znRecord = propertyStore.get(path, stat, AccessOption.PERSISTENT);
  if (LOG.isTraceEnabled()) {
    LOG.trace("Stat => " + stat);
  }
  long checkpoint = znRecord != null ? znRecord.getLongField(ZNODE_FIELD, -1) : -1;
  LOG.debug("Checkpoint value = " + checkpoint);
  return checkpoint;
}
 

/**
 * 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();
  }
}
 

/**
 * Set the runtime context of a single workflow
 * @param manager a connection to Helix
 * @param workflow the name of the workflow
 * @param workflowContext the up-to-date {@link WorkflowContext} for the workflow
 */
protected static void setWorkflowContext(HelixManager manager, String workflow,
    WorkflowContext workflowContext) {
  manager.getHelixPropertyStore().set(
      Joiner.on("/").join(TaskConstants.REBALANCER_CONTEXT_ROOT, workflow, CONTEXT_NODE),
      workflowContext.getRecord(), AccessOption.PERSISTENT);
}
 

private void ensurePropertyStoreEntryExistsForHighLevelConsumer(String realtimeTableName) {
  String propertyStorePath = ZKMetadataProvider.constructPropertyStorePathForResource(realtimeTableName);
  if (!_propertyStore.exists(propertyStorePath, AccessOption.PERSISTENT)) {
    LOGGER.info("Creating property store entry for HLC table: {}", realtimeTableName);
    _propertyStore.create(propertyStorePath, new ZNRecord(realtimeTableName), AccessOption.PERSISTENT);
  }
}