Examples with DataUpdater org.apache.helix.zookeeper.zkclient.DataUpdater used on opensource projects

Example 6 with DataUpdater

use of org.apache.helix.zookeeper.zkclient.DataUpdater in project helix by apache.

the class VirtualTopologyGroupService method createInstanceConfigUpdater.

/**
 * Create updater for instance config for async update.
 * @param clusterName cluster name of the instances.
 * @param assignment virtual group assignment.
 * @return a map from instance zkPath to its {@link DataUpdater} to update.
 */
@VisibleForTesting
static Map<String, DataUpdater<ZNRecord>> createInstanceConfigUpdater(String clusterName, Map<String, Set<String>> assignment) {
    Map<String, DataUpdater<ZNRecord>> updaters = new HashMap<>();
    for (Map.Entry<String, Set<String>> entry : assignment.entrySet()) {
        String virtualGroup = entry.getKey();
        for (String instanceName : entry.getValue()) {
            String path = PropertyPathBuilder.instanceConfig(clusterName, instanceName);
            updaters.put(path, currentData -> {
                InstanceConfig instanceConfig = new InstanceConfig(currentData);
                Map<String, String> domainMap = instanceConfig.getDomainAsMap();
                domainMap.put(VirtualTopologyGroupConstants.VIRTUAL_FAULT_ZONE_TYPE, virtualGroup);
                instanceConfig.setDomain(domainMap);
                return instanceConfig.getRecord();
            });
        }
    }
    return updaters;
}

Example 7 with DataUpdater

use of org.apache.helix.zookeeper.zkclient.DataUpdater in project helix by apache.

the class TestZkCacheAsyncOpSingleThread method testHappyPathSelfOpZkCacheBaseDataAccessor.

@Test
public void testHappyPathSelfOpZkCacheBaseDataAccessor() {
    String className = TestHelper.getTestClassName();
    String methodName = TestHelper.getTestMethodName();
    String clusterName = className + "_" + methodName;
    System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));
    // init zkCacheDataAccessor
    String curStatePath = PropertyPathBuilder.instanceCurrentState(clusterName, "localhost_8901");
    String extViewPath = PropertyPathBuilder.externalView(clusterName);
    ZkBaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<>(_gZkClient);
    baseAccessor.create(curStatePath, null, AccessOption.PERSISTENT);
    List<String> zkCacheInitPaths = Arrays.asList(curStatePath, extViewPath);
    ZkCacheBaseDataAccessor<ZNRecord> accessor = new ZkCacheBaseDataAccessor<>(baseAccessor, null, null, zkCacheInitPaths);
    // TestHelper.printCache(accessor._zkCache._cache);
    boolean ret = TestHelper.verifyZkCache(zkCacheInitPaths, accessor._zkCache._cache, _gZkClient, true);
    Assert.assertTrue(ret, "zkCache doesn't match data on Zk");
    // create 10 current states using this accessor
    List<String> paths = new ArrayList<>();
    List<ZNRecord> records = new ArrayList<>();
    for (int i = 0; i < 10; i++) {
        String path = PropertyPathBuilder.instanceCurrentState(clusterName, "localhost_8901", "session_0", "TestDB" + i);
        ZNRecord record = new ZNRecord("TestDB" + i);
        paths.add(path);
        records.add(record);
    }
    boolean[] success = accessor.createChildren(paths, records, AccessOption.PERSISTENT);
    for (int i = 0; i < 10; i++) {
        Assert.assertTrue(success[i], "Should succeed in create: " + paths.get(i));
    }
    // verify cache
    // TestHelper.printCache(accessor._zkCache._cache);
    ret = TestHelper.verifyZkCache(zkCacheInitPaths, accessor._zkCache._cache, _gZkClient, false);
    Assert.assertTrue(ret, "zkCache doesn't match data on Zk");
    // update each current state 10 times by this accessor
    List<DataUpdater<ZNRecord>> updaters = new ArrayList<>();
    for (int j = 0; j < 10; j++) {
        paths.clear();
        updaters.clear();
        for (int i = 0; i < 10; i++) {
            String path = curStatePath + "/session_0/TestDB" + i;
            ZNRecord newRecord = new ZNRecord("TestDB" + i);
            newRecord.setSimpleField("" + j, "" + j);
            DataUpdater<ZNRecord> updater = new ZNRecordUpdater(newRecord);
            paths.add(path);
            updaters.add(updater);
        }
        success = accessor.updateChildren(paths, updaters, AccessOption.PERSISTENT);
        for (int i = 0; i < 10; i++) {
            Assert.assertTrue(success[i], "Should succeed in update: " + paths.get(i));
        }
    }
    // verify cache
    // TestHelper.printCache(accessor._zkCache._cache);
    ret = TestHelper.verifyZkCache(zkCacheInitPaths, zkCacheInitPaths, accessor._zkCache._cache, _gZkClient, true);
    // ret = TestHelper.verifyZkCache(zkCacheInitPaths, accessor, _gZkClient, true);
    // System.out.println("ret: " + ret);
    Assert.assertTrue(ret, "zkCache doesn't match data on Zk");
    // set 10 external views 10 times by this accessor
    paths.clear();
    records.clear();
    for (int j = 0; j < 10; j++) {
        for (int i = 0; i < 10; i++) {
            String path = PropertyPathBuilder.externalView(clusterName, "TestDB" + i);
            ZNRecord record = new ZNRecord("TestDB" + i);
            record.setSimpleField("setKey", "" + j);
            paths.add(path);
            records.add(record);
        }
        success = accessor.setChildren(paths, records, AccessOption.PERSISTENT);
        for (int i = 0; i < 10; i++) {
            Assert.assertTrue(success[i], "Should succeed in set: " + paths.get(i));
        }
    }
    // verify cache
    // TestHelper.printCache(accessor._zkCache._cache);
    ret = TestHelper.verifyZkCache(zkCacheInitPaths, accessor._zkCache._cache, _gZkClient, true);
    // System.out.println("ret: " + ret);
    Assert.assertTrue(ret, "zkCache doesn't match data on Zk");
    // get 10 external views
    paths.clear();
    records.clear();
    for (int i = 0; i < 10; i++) {
        String path = extViewPath + "/TestDB" + i;
        paths.add(path);
    }
    records = accessor.get(paths, null, 0, true);
    for (int i = 0; i < 10; i++) {
        Assert.assertEquals(records.get(i).getId(), "TestDB" + i);
    }
    // getChildren
    records.clear();
    records = accessor.getChildren(extViewPath, null, 0, 0, 0);
    for (int i = 0; i < 10; i++) {
        Assert.assertEquals(records.get(i).getId(), "TestDB" + i);
    }
    // // exists
    paths.clear();
    for (int i = 0; i < 10; i++) {
        String path = curStatePath + "/session_0/TestDB" + i;
        // // PropertyPathConfig.getPath(PropertyType.CURRENTSTATES,
        // // clusterName,
        // // "localhost_8901",
        // // "session_0",
        // // "TestDB" + i);
        paths.add(path);
    }
    success = accessor.exists(paths, 0);
    for (int i = 0; i < 10; i++) {
        Assert.assertTrue(success[i], "Should exits: " + paths.get(i));
    }
    deleteCluster(clusterName);
    System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}

Example 8 with DataUpdater

use of org.apache.helix.zookeeper.zkclient.DataUpdater in project helix by apache.

the class TestZkCacheAsyncOpSingleThread method testHappyPathExtOpZkCacheBaseDataAccessor.

@Test
public void testHappyPathExtOpZkCacheBaseDataAccessor() throws Exception {
    String className = TestHelper.getTestClassName();
    String methodName = TestHelper.getTestMethodName();
    String clusterName = className + "_" + methodName;
    System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));
    // init external base data accessor
    HelixZkClient extZkclient = SharedZkClientFactory.getInstance().buildZkClient(new HelixZkClient.ZkConnectionConfig(ZK_ADDR));
    extZkclient.setZkSerializer(new ZNRecordSerializer());
    ZkBaseDataAccessor<ZNRecord> extBaseAccessor = new ZkBaseDataAccessor<>(extZkclient);
    // init zkCacheBaseDataAccessor
    String curStatePath = PropertyPathBuilder.instanceCurrentState(clusterName, "localhost_8901");
    String extViewPath = PropertyPathBuilder.externalView(clusterName);
    ZkBaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<>(_gZkClient);
    extBaseAccessor.create(curStatePath, null, AccessOption.PERSISTENT);
    List<String> zkCacheInitPaths = Arrays.asList(curStatePath, extViewPath);
    ZkCacheBaseDataAccessor<ZNRecord> accessor = new ZkCacheBaseDataAccessor<>(baseAccessor, null, null, zkCacheInitPaths);
    // TestHelper.printCache(accessor._zkCache);
    boolean ret = TestHelper.verifyZkCache(zkCacheInitPaths, accessor._zkCache._cache, _gZkClient, true);
    Assert.assertTrue(ret, "zkCache doesn't match data on Zk");
    // create 10 current states using external base accessor
    List<String> paths = new ArrayList<>();
    List<ZNRecord> records = new ArrayList<>();
    for (int i = 0; i < 10; i++) {
        String path = PropertyPathBuilder.instanceCurrentState(clusterName, "localhost_8901", "session_0", "TestDB" + i);
        ZNRecord record = new ZNRecord("TestDB" + i);
        paths.add(path);
        records.add(record);
    }
    boolean[] success = extBaseAccessor.createChildren(paths, records, AccessOption.PERSISTENT);
    for (int i = 0; i < 10; i++) {
        Assert.assertTrue(success[i], "Should succeed in create: " + paths.get(i));
    }
    // verify wtCache
    for (int i = 0; i < 10; i++) {
        // TestHelper.printCache(accessor._zkCache);
        ret = TestHelper.verifyZkCache(zkCacheInitPaths, accessor._zkCache._cache, _gZkClient, true);
        if (ret)
            break;
        Thread.sleep(100);
    }
    // System.out.println("ret: " + ret);
    Assert.assertTrue(ret, "zkCache doesn't match data on Zk");
    // update each current state 10 times by external base accessor
    List<DataUpdater<ZNRecord>> updaters = new ArrayList<>();
    for (int j = 0; j < 10; j++) {
        paths.clear();
        updaters.clear();
        for (int i = 0; i < 10; i++) {
            String path = curStatePath + "/session_0/TestDB" + i;
            ZNRecord newRecord = new ZNRecord("TestDB" + i);
            newRecord.setSimpleField("" + j, "" + j);
            DataUpdater<ZNRecord> updater = new ZNRecordUpdater(newRecord);
            paths.add(path);
            updaters.add(updater);
        }
        success = extBaseAccessor.updateChildren(paths, updaters, AccessOption.PERSISTENT);
        for (int i = 0; i < 10; i++) {
            Assert.assertTrue(success[i], "Should succeed in update: " + paths.get(i));
        }
    }
    // wait zkEventThread update zkCache
    Thread.sleep(100);
    // verify cache
    // TestHelper.printCache(accessor._zkCache);
    ret = TestHelper.verifyZkCache(zkCacheInitPaths, accessor._zkCache._cache, _gZkClient, true);
    Assert.assertTrue(ret, "zkCache doesn't match data on Zk");
    // set 10 external views by external accessor
    paths.clear();
    records.clear();
    for (int i = 0; i < 10; i++) {
        String path = PropertyPathBuilder.externalView(clusterName, "TestDB" + i);
        ZNRecord record = new ZNRecord("TestDB" + i);
        paths.add(path);
        records.add(record);
    }
    success = extBaseAccessor.setChildren(paths, records, AccessOption.PERSISTENT);
    for (int i = 0; i < 10; i++) {
        Assert.assertTrue(success[i], "Should succeed in set: " + paths.get(i));
    }
    // wait zkEventThread update zkCache
    Thread.sleep(100);
    // verify cache
    // TestHelper.printCache(accessor._zkCache._cache);
    ret = TestHelper.verifyZkCache(zkCacheInitPaths, accessor._zkCache._cache, _gZkClient, true);
    Assert.assertTrue(ret, "zkCache doesn't match data on Zk");
    // remove 10 external views by external accessor
    paths.clear();
    for (int i = 0; i < 10; i++) {
        String path = PropertyPathBuilder.externalView(clusterName, "TestDB" + i);
        paths.add(path);
    }
    success = extBaseAccessor.remove(paths, 0);
    for (int i = 0; i < 10; i++) {
        Assert.assertTrue(success[i], "Should succeed in remove: " + paths.get(i));
    }
    // wait zkEventThread update zkCache
    Thread.sleep(100);
    // verify cache
    // TestHelper.printCache(accessor._zkCache._cache);
    ret = TestHelper.verifyZkCache(zkCacheInitPaths, accessor._zkCache._cache, _gZkClient, true);
    Assert.assertTrue(ret, "zkCache doesn't match data on Zk");
    // clean up
    extZkclient.close();
    deleteCluster(clusterName);
    System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}

Example 9 with DataUpdater

use of org.apache.helix.zookeeper.zkclient.DataUpdater in project helix by apache.

the class TaskDriver method enqueueJobs.

/**
 * Batch add jobs to queues that garantee
 * @param queue
 * @param jobs
 * @param jobBuilders
 */
public void enqueueJobs(final String queue, final List<String> jobs, final List<JobConfig.Builder> jobBuilders) {
    // Get the job queue config and capacity
    WorkflowConfig workflowConfig = TaskUtil.getWorkflowConfig(_accessor, queue);
    if (workflowConfig == null) {
        throw new IllegalArgumentException("Queue " + queue + " config does not yet exist!");
    }
    if (workflowConfig.isTerminable()) {
        throw new IllegalArgumentException(queue + " is not a queue!");
    }
    final int capacity = workflowConfig.getCapacity();
    int queueSize = workflowConfig.getJobDag().size();
    if (capacity > 0 && queueSize >= capacity) {
        // if queue is full, Helix will try to clean up the expired job to free more space.
        WorkflowContext workflowContext = TaskUtil.getWorkflowContext(_propertyStore, queue);
        if (workflowContext != null) {
            Set<String> expiredJobs = TaskUtil.getExpiredJobs(_accessor, _propertyStore, workflowConfig, workflowContext);
            if (!TaskUtil.removeJobsFromWorkflow(_accessor, _propertyStore, queue, expiredJobs, true)) {
                LOG.warn("Failed to clean up expired and completed jobs from queue {}", queue);
            }
        }
        workflowConfig = TaskUtil.getWorkflowConfig(_accessor, queue);
        if (workflowConfig.getJobDag().size() >= capacity) {
            throw new HelixException(String.format("Failed to enqueue job, queue %s is full.", queue));
        }
    }
    // Fail the operation if adding new jobs will cause the queue to reach its capacity limit
    workflowConfig = TaskUtil.getWorkflowConfig(_accessor, queue);
    if (workflowConfig.getJobDag().size() + jobs.size() >= capacity) {
        throw new IllegalStateException(String.format("Queue %s already reaches its max capacity %d, failed to add %s", queue, capacity, jobs.toString()));
    }
    validateZKNodeLimitation(1);
    final List<JobConfig> jobConfigs = new ArrayList<>();
    final List<String> namespacedJobNames = new ArrayList<>();
    final List<String> jobTypeList = new ArrayList<>();
    try {
        for (int i = 0; i < jobBuilders.size(); i++) {
            // Create the job to ensure that it validates
            JobConfig jobConfig = jobBuilders.get(i).setWorkflow(queue).build();
            String namespacedJobName = TaskUtil.getNamespacedJobName(queue, jobs.get(i));
            // add job config first.
            addJobConfig(namespacedJobName, jobConfig);
            jobConfigs.add(jobConfig);
            namespacedJobNames.add(namespacedJobName);
            jobTypeList.add(jobConfig.getJobType());
        }
    } catch (HelixException e) {
        LOG.error("Failed to add job configs {}. Remove them all!", jobs.toString());
        for (String job : jobs) {
            String namespacedJobName = TaskUtil.getNamespacedJobName(queue, job);
            TaskUtil.removeJobConfig(_accessor, namespacedJobName);
        }
    }
    // update the job dag to append the job to the end of the queue.
    DataUpdater<ZNRecord> updater = currentData -> {
        if (currentData == null) {
            // In this case, we cannot proceed and must alert the user
            throw new HelixException(String.format("enqueueJobs DataUpdater: JobQueue %s config is not found!", queue));
        }
        // Add the node to the existing DAG
        JobDag jobDag = JobDag.fromJson(currentData.getSimpleField(WorkflowConfig.WorkflowConfigProperty.Dag.name()));
        Set<String> allNodes = jobDag.getAllNodes();
        if (capacity > 0 && allNodes.size() + jobConfigs.size() >= capacity) {
            // same time and cause overcapacity queue
            for (String job : jobs) {
                String namespacedJobName = TaskUtil.getNamespacedJobName(queue, job);
                TaskUtil.removeJobConfig(_accessor, namespacedJobName);
            }
            throw new IllegalStateException(String.format("Queue %s already reaches its max capacity %d, failed to add %s", queue, capacity, jobs.toString()));
        }
        String lastNodeName = null;
        for (int i = 0; i < namespacedJobNames.size(); i++) {
            String namespacedJobName = namespacedJobNames.get(i);
            if (allNodes.contains(namespacedJobName)) {
                throw new IllegalStateException(String.format("Could not add to queue %s, job %s already exists", queue, jobs.get(i)));
            }
            jobDag.addNode(namespacedJobName);
            // Add the node to the end of the queue
            String candidate = null;
            if (lastNodeName == null) {
                for (String node : allNodes) {
                    if (!node.equals(namespacedJobName) && jobDag.getDirectChildren(node).isEmpty()) {
                        candidate = node;
                        break;
                    }
                }
            } else {
                candidate = lastNodeName;
            }
            if (candidate != null) {
                jobDag.addParentToChild(candidate, namespacedJobName);
                lastNodeName = namespacedJobName;
            }
        }
        // Add job type if job type is not null
        Map<String, String> jobTypes = currentData.getMapField(WorkflowConfig.WorkflowConfigProperty.JobTypes.name());
        for (String jobType : jobTypeList) {
            if (jobType != null) {
                if (jobTypes == null) {
                    jobTypes = new HashMap<>();
                }
                jobTypes.put(queue, jobType);
            }
        }
        if (jobTypes != null) {
            currentData.setMapField(WorkflowConfig.WorkflowConfigProperty.JobTypes.name(), jobTypes);
        }
        // Save the updated DAG
        try {
            currentData.setSimpleField(WorkflowConfig.WorkflowConfigProperty.Dag.name(), jobDag.toJson());
        } catch (Exception e) {
            throw new IllegalStateException(String.format("Could not add jobs %s to queue %s", jobs.toString(), queue), e);
        }
        return currentData;
    };
    String path = _accessor.keyBuilder().resourceConfig(queue).getPath();
    boolean status = _accessor.getBaseDataAccessor().update(path, updater, AccessOption.PERSISTENT);
    if (!status) {
        LOG.error("Failed to update WorkflowConfig, remove all jobs {}", jobs.toString());
        for (String job : jobs) {
            TaskUtil.removeJobConfig(_accessor, job);
        }
        throw new HelixException("Failed to enqueue job");
    }
}

Example 10 with DataUpdater

use of org.apache.helix.zookeeper.zkclient.DataUpdater in project helix by apache.

the class HelixTaskExecutor method updateMessageState.

private void updateMessageState(Collection<Message> msgsToBeUpdated, HelixDataAccessor accessor, String instanceName) {
    if (msgsToBeUpdated.isEmpty()) {
        return;
    }
    Builder keyBuilder = accessor.keyBuilder();
    List<Message> updateMsgs = new ArrayList<>();
    List<String> updateMsgPaths = new ArrayList<>();
    List<DataUpdater<ZNRecord>> updaters = new ArrayList<>();
    for (Message msg : msgsToBeUpdated) {
        updateMsgs.add(msg);
        updateMsgPaths.add(msg.getKey(keyBuilder, instanceName).getPath());
        /**
         * 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(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();
        });
    }
    boolean[] updateResults = accessor.updateChildren(updateMsgPaths, updaters, AccessOption.PERSISTENT);
    boolean isMessageUpdatedAsNew = false;
    // This is to avoid inconsistent cache.
    for (int i = 0; i < updateMsgs.size(); i++) {
        Message msg = updateMsgs.get(i);
        if (msg.getMsgState().equals(MessageState.NEW)) {
            // If a message is updated as NEW state, then we might need to process it again soon.
            isMessageUpdatedAsNew = true;
        // And it shall not be treated as a known messages.
        } else {
            _knownMessageIds.add(msg.getId());
            if (!updateResults[i]) {
                // TODO: If the message update fails, maybe we shall not treat the message as a known
                // TODO: message. We shall apply more strict check and retry the update.
                LOG.error("Failed to update the message {}.", msg.getMsgId());
            }
        }
    }
    if (isMessageUpdatedAsNew) {
        // Sending a NO-OP message to trigger another message callback to re-process the messages
        // that are updated as NEW state.
        sendNopMessage(accessor, instanceName);
    }
}