Examples with LoadDistributionService org.apache.nifi.loading.LoadDistributionService used on opensource projects

Example 1 with LoadDistributionService

use of org.apache.nifi.loading.LoadDistributionService in project nifi by apache.

the class DistributeLoad method createWeightedList.

@OnScheduled
public void createWeightedList(final ProcessContext context) {
    final Map<Integer, Integer> weightings = new LinkedHashMap<>();
    String distStrat = context.getProperty(DISTRIBUTION_STRATEGY).getValue();
    if (distStrat.equals(STRATEGY_LOAD_DISTRIBUTION_SERVICE)) {
        String hostNamesValue = context.getProperty(HOSTNAMES).getValue();
        String[] hostNames = hostNamesValue.split("(?:,+|;+|\\s+)");
        Set<String> hostNameSet = new HashSet<>();
        for (String hostName : hostNames) {
            if (StringUtils.isNotBlank(hostName)) {
                hostNameSet.add(hostName);
            }
        }
        LoadDistributionService svc = context.getProperty(LOAD_DISTRIBUTION_SERVICE_TEMPLATE).asControllerService(LoadDistributionService.class);
        myListener = new LoadDistributionListener() {

            @Override
            public void update(Map<String, Integer> loadInfo) {
                for (Relationship rel : relationshipsRef.get()) {
                    String hostname = rel.getDescription();
                    Integer weight = 1;
                    if (loadInfo.containsKey(hostname)) {
                        weight = loadInfo.get(hostname);
                    }
                    weightings.put(Integer.decode(rel.getName()), weight);
                }
                updateWeightedRelationships(weightings);
            }
        };
        Map<String, Integer> loadInfo = svc.getLoadDistribution(hostNameSet, myListener);
        for (Relationship rel : relationshipsRef.get()) {
            String hostname = rel.getDescription();
            Integer weight = 1;
            if (loadInfo.containsKey(hostname)) {
                weight = loadInfo.get(hostname);
            }
            weightings.put(Integer.decode(rel.getName()), weight);
        }
    } else {
        final int numRelationships = context.getProperty(NUM_RELATIONSHIPS).asInteger();
        for (int i = 1; i <= numRelationships; i++) {
            weightings.put(i, 1);
        }
        for (final PropertyDescriptor propDesc : context.getProperties().keySet()) {
            if (!this.properties.contains(propDesc)) {
                final int relationship = Integer.parseInt(propDesc.getName());
                final int weighting = context.getProperty(propDesc).asInteger();
                weightings.put(relationship, weighting);
            }
        }
    }
    updateWeightedRelationships(weightings);
}