Examples with AffinityRule com.datatorrent.api.AffinityRule used on opensource projects

Example 6 with AffinityRule

use of com.datatorrent.api.AffinityRule in project apex-core by apache.

the class AffinityRulesTest method testAntiAffinityInOperators.

@Test
public void testAntiAffinityInOperators() {
    LogicalPlan dag = new LogicalPlan();
    dag.getAttributes().put(com.datatorrent.api.Context.DAGContext.APPLICATION_PATH, testMeta.getAbsolutePath());
    dag.setAttribute(OperatorContext.STORAGE_AGENT, new MemoryStorageAgent());
    GenericTestOperator o1 = dag.addOperator("O1", GenericTestOperator.class);
    dag.setOperatorAttribute(o1, OperatorContext.MEMORY_MB, 256);
    GenericTestOperator o2 = dag.addOperator("O2", GenericTestOperator.class);
    dag.setOperatorAttribute(o2, OperatorContext.MEMORY_MB, 256);
    dag.getMeta(o1).getAttributes().put(OperatorContext.LOCALITY_HOST, "host1");
    AffinityRulesSet ruleSet = new AffinityRulesSet();
    List<AffinityRule> rules = new ArrayList<>();
    ruleSet.setAffinityRules(rules);
    AffinityRule rule1 = new AffinityRule(Type.ANTI_AFFINITY, Locality.NODE_LOCAL, false, "O1", "O2");
    rules.add(rule1);
    dag.setAttribute(DAGContext.AFFINITY_RULES_SET, ruleSet);
    // .setLocality(Locality.NODE_LOCAL);
    dag.addStream("o1_outport1", o1.outport1, o2.inport1);
    StreamingContainerManager scm = new StreamingContainerManager(dag);
    ResourceRequestHandler rr = new ResourceRequestHandler();
    int containerMem = 1000;
    Map<String, NodeReport> nodeReports = Maps.newHashMap();
    NodeReport nr = BuilderUtils.newNodeReport(BuilderUtils.newNodeId("host1", 0), NodeState.RUNNING, "httpAddress", "rackName", BuilderUtils.newResource(0, 0), BuilderUtils.newResource(containerMem * 2, 2), 0, null, 0);
    nodeReports.put(nr.getNodeId().getHost(), nr);
    nr = BuilderUtils.newNodeReport(BuilderUtils.newNodeId("host2", 0), NodeState.RUNNING, "httpAddress", "rackName", BuilderUtils.newResource(0, 0), BuilderUtils.newResource(containerMem * 2, 2), 0, null, 0);
    nodeReports.put(nr.getNodeId().getHost(), nr);
    // set resources
    rr.updateNodeReports(Lists.newArrayList(nodeReports.values()));
    for (ContainerStartRequest csr : scm.containerStartRequests) {
        String host = rr.getHost(csr, true);
        csr.container.host = host;
        if (csr.container.getOperators().get(0).getName().equals("O1")) {
            Assert.assertEquals("Hosts set to host1 for Operator O1", "host1", host);
        }
        if (csr.container.getOperators().get(0).getName().equals("O2")) {
            Assert.assertEquals("Hosts set to host2 for Operator O2", "host2", host);
        }
    }
}