Examples with Assignment org.apache.pulsar.functions.proto.Function.Assignment used on opensource projects

Example 6 with Assignment

use of org.apache.pulsar.functions.proto.Function.Assignment in project incubator-pulsar by apache.

the class SchedulerManager method invokeScheduler.

private void invokeScheduler() {
    List<String> currentMembership = this.membershipManager.getCurrentMembership().stream().map(workerInfo -> workerInfo.getWorkerId()).collect(Collectors.toList());
    List<FunctionMetaData> allFunctions = this.functionMetaDataManager.getAllFunctionMetaData();
    Map<String, Function.Instance> allInstances = computeAllInstances(allFunctions);
    Map<String, Map<String, Assignment>> workerIdToAssignments = this.functionRuntimeManager.getCurrentAssignments();
    // delete assignments of functions and instances that don't exist anymore
    Iterator<Map.Entry<String, Map<String, Assignment>>> it = workerIdToAssignments.entrySet().iterator();
    while (it.hasNext()) {
        Map.Entry<String, Map<String, Assignment>> workerIdToAssignmentEntry = it.next();
        Map<String, Assignment> functionMap = workerIdToAssignmentEntry.getValue();
        // remove instances that don't exist anymore
        functionMap.entrySet().removeIf(entry -> {
            String fullyQualifiedInstanceId = entry.getKey();
            return !allInstances.containsKey(fullyQualifiedInstanceId);
        });
        // update assignment instances in case attributes of a function gets updated
        for (Map.Entry<String, Assignment> entry : functionMap.entrySet()) {
            String fullyQualifiedInstanceId = entry.getKey();
            Assignment assignment = entry.getValue();
            Function.Instance instance = allInstances.get(fullyQualifiedInstanceId);
            if (!assignment.getInstance().equals(instance)) {
                functionMap.put(fullyQualifiedInstanceId, assignment.toBuilder().setInstance(instance).build());
            }
        }
        if (functionMap.isEmpty()) {
            it.remove();
        }
    }
    List<Assignment> currentAssignments = workerIdToAssignments.entrySet().stream().flatMap(stringMapEntry -> stringMapEntry.getValue().values().stream()).collect(Collectors.toList());
    List<Function.Instance> needsAssignment = this.getUnassignedFunctionInstances(workerIdToAssignments, allInstances);
    List<Assignment> assignments = this.scheduler.schedule(needsAssignment, currentAssignments, currentMembership);
    log.debug("New assignments computed: {}", assignments);
    long assignmentVersion = this.functionRuntimeManager.getCurrentAssignmentVersion() + 1;
    Request.AssignmentsUpdate assignmentsUpdate = Request.AssignmentsUpdate.newBuilder().setVersion(assignmentVersion).addAllAssignments(assignments).build();
    CompletableFuture<MessageId> messageIdCompletableFuture = producer.sendAsync(assignmentsUpdate.toByteArray());
    try {
        messageIdCompletableFuture.get();
    } catch (InterruptedException | ExecutionException e) {
        log.error("Failed to send assignment update", e);
        throw new RuntimeException(e);
    }
    // wait for assignment update to go throw the pipeline
    int retries = 0;
    while (this.functionRuntimeManager.getCurrentAssignmentVersion() < assignmentVersion) {
        if (retries >= this.workerConfig.getAssignmentWriteMaxRetries()) {
            log.warn("Max number of retries reached for waiting for assignment to propagate. Will continue now.");
            break;
        }
        log.info("Waiting for assignments to propagate...");
        try {
            Thread.sleep(500);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        retries++;
    }
}