Java Code Examples for com.datatorrent.stram.plan.logical.LogicalPlan#addOperator()

@Test
public void testContainerSizeWithPartitioning()
{
  LogicalPlan dag = new LogicalPlan();
  dag.setAttribute(OperatorContext.STORAGE_AGENT, new StramTestSupport.MemoryStorageAgent());

  GenericTestOperator o1 = dag.addOperator("o1", GenericTestOperator.class);
  GenericTestOperator o2 = dag.addOperator("o2", GenericTestOperator.class);
  dag.setOperatorAttribute(o1, OperatorContext.PARTITIONER, new StatelessPartitioner<GenericTestOperator>(3));
  dag.setOperatorAttribute(o2, OperatorContext.PARTITIONER, new StatelessPartitioner<GenericTestOperator>(2));

  dag.addStream("o1.outport1", o1.outport1, o2.inport1);
  dag.setAttribute(LogicalPlan.CONTAINERS_MAX_COUNT, 10);
  PhysicalPlan plan = new PhysicalPlan(dag, new TestPlanContext());
  Assert.assertEquals("number of containers", 5, plan.getContainers().size());
  PTContainer container;
  for (int i = 0; i < 5; i++) {
    container = plan.getContainers().get(i);
    if (container.getOperators().size() == 1) {
      Assert.assertEquals("container memory is 1536 for container :" + container, 1536, container.getRequiredMemoryMB());
    }
    if (container.getOperators().size() == 2) {
      Assert.assertEquals("container memory is 2048 for container :" + container, 2048, container.getRequiredMemoryMB());
    }
  }
}
 

@Test
public void validatePositiveOiOiO()
{
  logger.info("Checking the logic for sanity checking of OiO");

  LogicalPlan plan = new LogicalPlan();
  ThreadIdValidatingInputOperator inputOperator = plan.addOperator("inputOperator", new ThreadIdValidatingInputOperator());
  ThreadIdValidatingGenericIntermediateOperator intermediateOperator = plan.addOperator("intermediateOperator", new ThreadIdValidatingGenericIntermediateOperator());
  ThreadIdValidatingOutputOperator outputOperator = plan.addOperator("outputOperator", new ThreadIdValidatingOutputOperator());

  plan.addStream("OiO1", inputOperator.output, intermediateOperator.input).setLocality(Locality.THREAD_LOCAL);
  plan.addStream("OiO2", intermediateOperator.output, outputOperator.input).setLocality(Locality.THREAD_LOCAL);

  try {
    plan.validate();
    Assert.assertTrue("OiOiO validation", true);
  } catch (ConstraintViolationException ex) {
    Assert.fail("OiOiO validation");
  }
}
 

@Test
public void testCommitted() throws IOException, ClassNotFoundException
{
  LogicalPlan lp = new LogicalPlan();
  String workingDir = new File("target/testCommitted").getAbsolutePath();
  lp.setAttribute(Context.OperatorContext.STORAGE_AGENT, new AsyncFSStorageAgent(workingDir, null));
  lp.setAttribute(DAGContext.CHECKPOINT_WINDOW_COUNT, 1);
  String opName = "CommitAwareOperatorTestCommit";
  lp.addOperator(opName, new CommitAwareOperator());

  StramLocalCluster lc = new StramLocalCluster(lp);
  lc.run(5000);

  /* this is not foolproof but some insurance is better than nothing */
  Assert.assertTrue("No Committed Windows", committedWindowIds.contains(opName));
}
 

@Test
public void testParallelPartitionForSlidingWindow()
{
  LogicalPlan dag = new LogicalPlan();
  dag.setAttribute(OperatorContext.STORAGE_AGENT, new StramTestSupport.MemoryStorageAgent());

  GenericTestOperator o1 = dag.addOperator("o1", GenericTestOperator.class);
  GenericTestOperator o2 = dag.addOperator("o2", GenericTestOperator.class);
  GenericTestOperator o3 = dag.addOperator("o3", GenericTestOperator.class);
  dag.setOperatorAttribute(o1, OperatorContext.SLIDE_BY_WINDOW_COUNT, 2);
  dag.setOperatorAttribute(o1, OperatorContext.PARTITIONER, new StatelessPartitioner<>(2));
  dag.setInputPortAttribute(o2.inport1, PortContext.PARTITION_PARALLEL, true);
  dag.setOperatorAttribute(o1, OperatorContext.APPLICATION_WINDOW_COUNT, 4);

  dag.addStream("o1.outport1", o1.outport1, o2.inport1);
  dag.addStream("o2.outport1", o2.outport1, o3.inport1);
  PhysicalPlan plan = new PhysicalPlan(dag, new TestPlanContext());
  Assert.assertEquals("number of containers", 7, plan.getContainers().size());
}
 

@Test
public void validatePositiveOiOiOExtendeddiamond()
{
  logger.info("Checking the logic for sanity checking of OiO");

  LogicalPlan plan = new LogicalPlan();
  ThreadIdValidatingInputOperator inputOperator = plan.addOperator("inputOperator", new ThreadIdValidatingInputOperator());
  ThreadIdValidatingGenericIntermediateOperator intermediateOperator1 = plan.addOperator("intermediateOperator1", new ThreadIdValidatingGenericIntermediateOperator());
  ThreadIdValidatingGenericIntermediateOperator intermediateOperator2 = plan.addOperator("intermediateOperator2", new ThreadIdValidatingGenericIntermediateOperator());
  ThreadIdValidatingGenericIntermediateOperator intermediateOperator3 = plan.addOperator("intermediateOperator3", new ThreadIdValidatingGenericIntermediateOperator());
  ThreadIdValidatingGenericIntermediateOperator intermediateOperator4 = plan.addOperator("intermediateOperator4", new ThreadIdValidatingGenericIntermediateOperator());
  ThreadIdValidatingGenericOperatorWithTwoInputPorts outputOperator = plan.addOperator("outputOperator", new ThreadIdValidatingGenericOperatorWithTwoInputPorts());

  plan.addStream("OiOin", inputOperator.output, intermediateOperator1.input, intermediateOperator3.input).setLocality(Locality.THREAD_LOCAL);
  plan.addStream("OiOIntermediate1", intermediateOperator1.output, intermediateOperator2.input).setLocality(Locality.THREAD_LOCAL);
  plan.addStream("OiOIntermediate2", intermediateOperator3.output, intermediateOperator4.input).setLocality(Locality.THREAD_LOCAL);
  plan.addStream("OiOout1", intermediateOperator2.output, outputOperator.input).setLocality(Locality.THREAD_LOCAL);
  plan.addStream("OiOout2", intermediateOperator4.output, outputOperator.input2).setLocality(Locality.THREAD_LOCAL);

  try {
    plan.validate();
    Assert.assertTrue("OiOiO extended diamond validation", true);
  } catch (ConstraintViolationException ex) {
    Assert.fail("OIOIO extended diamond validation");
  }
}
 

/**
 * Module and Operator with same name is not allowed in a DAG, to prevent properties
 * conflict.
 */
@Test(expected = java.lang.IllegalArgumentException.class)
public void conflictingNamesWithOperator2()
{
  Configuration conf = new Configuration(false);
  LogicalPlanConfiguration lpc = new LogicalPlanConfiguration(conf);
  LogicalPlan dag = new LogicalPlan();
  Level2ModuleA module = dag.addModule("M1", new Level2ModuleA());
  DummyInputOperator in = dag.addOperator("M1", new DummyInputOperator());
  dag.addStream("s1", in.out, module.mIn);
  lpc.prepareDAG(dag, null, "ModuleApp");
  dag.validate();
}
 

@Test
public void testInlineMultipleInputs()
{

  LogicalPlan dag = new LogicalPlan();

  GenericTestOperator node1 = dag.addOperator("node1", GenericTestOperator.class);
  GenericTestOperator node2 = dag.addOperator("node2", GenericTestOperator.class);
  GenericTestOperator node3 = dag.addOperator("node3", GenericTestOperator.class);

  dag.addStream("n1Output1", node1.outport1, node3.inport1)
      .setLocality(Locality.CONTAINER_LOCAL);

  dag.addStream("n2Output1", node2.outport1, node3.inport2)
      .setLocality(Locality.CONTAINER_LOCAL);

  int maxContainers = 5;
  dag.setAttribute(LogicalPlan.CONTAINERS_MAX_COUNT, maxContainers);
  dag.setAttribute(OperatorContext.STORAGE_AGENT, new TestPlanContext());

  PhysicalPlan deployer = new PhysicalPlan(dag, new TestPlanContext());
  Assert.assertEquals("number of containers", 1, deployer.getContainers().size());

  PTOutput node1Out = deployer.getOperators(dag.getMeta(node1)).get(0).getOutputs().get(0);
  Assert.assertTrue("inline " + node1Out, node1Out.isDownStreamInline());

  // per current logic, different container is assigned to second input node
  PTOutput node2Out = deployer.getOperators(dag.getMeta(node2)).get(0).getOutputs().get(0);
  Assert.assertTrue("inline " + node2Out, node2Out.isDownStreamInline());

}
 

@Test
public void testContainerLocalWithVCores()
{
  LogicalPlan dag = new LogicalPlan();
  dag.getAttributes().put(com.datatorrent.api.Context.DAGContext.APPLICATION_PATH, new File("target", HostLocalTest.class.getName()).getAbsolutePath());
  dag.setAttribute(OperatorContext.STORAGE_AGENT, new MemoryStorageAgent());

  GenericTestOperator o1 = dag.addOperator("o1", GenericTestOperator.class);
  dag.getMeta(o1).getAttributes().put(OperatorContext.LOCALITY_HOST, "host2");

  GenericTestOperator partitioned = dag.addOperator("partitioned", GenericTestOperator.class);
  dag.addStream("o1_outport1", o1.outport1, partitioned.inport1).setLocality(Locality.CONTAINER_LOCAL);
  dag.setOperatorAttribute(o1,OperatorContext.MEMORY_MB,256);
  dag.setOperatorAttribute(o1,OperatorContext.VCORES,1);
  dag.setOperatorAttribute(partitioned,OperatorContext.VCORES,1);

  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()));
  Assert.assertEquals("number of containers is 1", 1, scm.containerStartRequests.size());
  for (ContainerStartRequest csr : scm.containerStartRequests) {
    String host = rr.getHost(csr, true);
    csr.container.host = host;
    Assert.assertEquals("number of vcores", 2, csr.container.getRequiredVCores());
    Assert.assertEquals("Hosts set to host2", "host2", host);
  }
}
 

@Test
public void testContainerLocal()
{
  LogicalPlan dag = new LogicalPlan();
  dag.getAttributes().put(com.datatorrent.api.Context.DAGContext.APPLICATION_PATH, new File("target", HostLocalTest.class.getName()).getAbsolutePath());
  dag.setAttribute(OperatorContext.STORAGE_AGENT, new MemoryStorageAgent());

  GenericTestOperator o1 = dag.addOperator("o1", GenericTestOperator.class);
  dag.getMeta(o1).getAttributes().put(OperatorContext.LOCALITY_HOST, "host2");

  GenericTestOperator partitioned = dag.addOperator("partitioned", GenericTestOperator.class);
  dag.addStream("o1_outport1", o1.outport1, partitioned.inport1).setLocality(Locality.CONTAINER_LOCAL);
  dag.setOperatorAttribute(o1, OperatorContext.MEMORY_MB, 256);
  dag.setOperatorAttribute(partitioned,OperatorContext.MEMORY_MB,256);

  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()));
  Assert.assertEquals("number of containers is 1", 1, scm.containerStartRequests.size());
  for (ContainerStartRequest csr : scm.containerStartRequests) {
    String host = rr.getHost(csr, true);
    csr.container.host = host;
    Assert.assertEquals("Hosts set to host2", "host2", host);
  }
}
 

@Test
public void testThreadLocal()
{
  LogicalPlan dag = new LogicalPlan();
  dag.getAttributes().put(com.datatorrent.api.Context.DAGContext.APPLICATION_PATH, new File("target", HostLocalTest.class.getName()).getAbsolutePath());
  dag.setAttribute(OperatorContext.STORAGE_AGENT, new MemoryStorageAgent());

  GenericTestOperator o1 = dag.addOperator("o1", GenericTestOperator.class);
  dag.getMeta(o1).getAttributes().put(OperatorContext.LOCALITY_HOST, "host2");

  GenericTestOperator partitioned = dag.addOperator("partitioned", GenericTestOperator.class);
  dag.addStream("o1_outport1", o1.outport1, partitioned.inport1).setLocality(Locality.THREAD_LOCAL);
  dag.setOperatorAttribute(o1,OperatorContext.MEMORY_MB,256);
  dag.setOperatorAttribute(partitioned,OperatorContext.MEMORY_MB,256);

  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()));
  Assert.assertEquals("number of containers is 1", 1, scm.containerStartRequests.size());
  for (ContainerStartRequest csr : scm.containerStartRequests) {
    String host = rr.getHost(csr, true);
    csr.container.host = host;
    Assert.assertEquals("Hosts set to host2", "host2", host);
  }
}
 

@Test
public void testContainerSize()
{
  LogicalPlan dag = new LogicalPlan();
  dag.setAttribute(OperatorContext.STORAGE_AGENT, new StramTestSupport.MemoryStorageAgent());

  GenericTestOperator o1 = dag.addOperator("o1", GenericTestOperator.class);
  GenericTestOperator o2 = dag.addOperator("o2", GenericTestOperator.class);
  GenericTestOperator o3 = dag.addOperator("o3", GenericTestOperator.class);
  dag.setOperatorAttribute(o1,OperatorContext.VCORES,1);
  dag.setOperatorAttribute(o2,OperatorContext.VCORES,2);

  dag.addStream("o1.outport1", o1.outport1, o2.inport1);
  dag.addStream("o2.outport1", o2.outport1, o3.inport1);

  dag.setOperatorAttribute(o2, OperatorContext.MEMORY_MB, 4000);
  dag.setAttribute(LogicalPlan.CONTAINERS_MAX_COUNT, 2);

  PhysicalPlan plan = new PhysicalPlan(dag, new TestPlanContext());

  Assert.assertEquals("number of containers", 2, plan.getContainers().size());
  Assert.assertEquals("memory container 1", 2560, plan.getContainers().get(0).getRequiredMemoryMB());
  Assert.assertEquals("vcores container 1", 1, plan.getContainers().get(0).getRequiredVCores());
  Assert.assertEquals("memory container 2", 4512, plan.getContainers().get(1).getRequiredMemoryMB());
  Assert.assertEquals("vcores container 2", 2, plan.getContainers().get(1).getRequiredVCores());
  Assert.assertEquals("number of operators in container 1", 2, plan.getContainers().get(0).getOperators().size());
}
 

@Test
  public void testInlineOperatorsRecovery() throws Exception
  {
    RecoverableInputOperator.initGenTuples();
    CollectorOperator.collection.clear();
    int maxTuples = 30;
    LogicalPlan dag = new LogicalPlan();
    String workingDir = new File("target/testOperatorRecovery").getAbsolutePath();
    AsyncFSStorageAgent asyncFSStorageAgent = new AsyncFSStorageAgent(workingDir, null);
    asyncFSStorageAgent.setSyncCheckpoint(true);
    dag.setAttribute(Context.OperatorContext.STORAGE_AGENT, asyncFSStorageAgent);
    //dag.getAttributes().get(DAG.HEARTBEAT_INTERVAL_MILLIS, 400);
    dag.getAttributes().put(LogicalPlan.CHECKPOINT_WINDOW_COUNT, 2);
    dag.getAttributes().put(LogicalPlan.STREAMING_WINDOW_SIZE_MILLIS, 300);
    dag.getAttributes().put(LogicalPlan.CONTAINERS_MAX_COUNT, 1);
    RecoverableInputOperator rip = dag.addOperator("LongGenerator", RecoverableInputOperator.class);
    rip.setMaximumTuples(maxTuples);
    rip.setSimulateFailure(true);

    CollectorOperator cm = dag.addOperator("LongCollector", CollectorOperator.class);
    cm.setSimulateFailure(true);
    dag.addStream("connection", rip.output, cm.input).setLocality(Locality.CONTAINER_LOCAL);

    StramLocalCluster lc = new StramLocalCluster(dag);
    lc.run();

//    for (Long l: collection) {
//      logger.debug(Codec.getStringWindowId(l));
//    }
    Assert.assertEquals("Generated Outputs", maxTuples, CollectorOperator.collection.size());
  }
 

@Test
public void testGenericOperatorPropertyCodec()
{
  LogicalPlan dag = new LogicalPlan();
  Map<Class<?>, Class<? extends StringCodec<?>>> codecs = new HashMap<>();
  codecs.put(GenericOperatorProperty.class, GenericOperatorProperty.GenericOperatorPropertyStringCodec.class);
  dag.setAttribute(com.datatorrent.api.Context.DAGContext.STRING_CODECS, codecs);
  dag.setAttribute(com.datatorrent.api.Context.OperatorContext.STORAGE_AGENT, new MemoryStorageAgent());

  GenericTestOperator o1 = dag.addOperator("o1", GenericTestOperator.class);
  OperatorMeta o1Meta = dag.getMeta(o1);

  TestPlanContext ctx = new TestPlanContext();
  PhysicalPlan plan = new PhysicalPlan(dag, ctx);
  ctx.deploy.clear();
  ctx.undeploy.clear();

  PlanModifier pm = new PlanModifier(plan);
  try {
    pm.setOperatorProperty(o1Meta.getName(), "genericOperatorProperty", "xyz");
    Assert.fail("validation error expected"); // cannot set properties on an operator that is already deployed.
  } catch (javax.validation.ValidationException e) {
    Assert.assertTrue(e.getMessage().contains(o1Meta.toString()));
  }

  GenericTestOperator newOperator = new GenericTestOperator();
  pm.addOperator("newOperator", newOperator);
  pm.setOperatorProperty("newOperator", "genericOperatorProperty", "xyz");
  Assert.assertEquals("", "xyz", newOperator.getGenericOperatorProperty().obtainString());
}
 

/**
 * Test that internally all stats listeners are handled through StatsListenerWithContext.
 * Following cases are tested
 *
 * Operator implementing StatsListener
 * Operator implementing StatsListenerWithContext
 * Operator with STATS_LISTENERS attribute set to StatsListener
 * Operator with STATS_LISTENERS attribute set to StatsListenerWithContext
 */
@Test
public void testStatsListenerContextWrappers()
{
  LogicalPlan dag = new LogicalPlan();
  dag.setAttribute(OperatorContext.STORAGE_AGENT, new StramTestSupport.MemoryStorageAgent());

  StatsListenerOperator o1 = dag.addOperator("o1", new StatsListenerOperator());
  GenericTestOperator o2 = dag.addOperator("o2", new GenericTestOperator());
  dag.setAttribute(o2, OperatorContext.STATS_LISTENERS,
      Lists.<StatsListener>newArrayList(mock(StatsListener.class)));

  GenericTestOperator o3 = dag.addOperator("o3", new GenericTestOperator());
  dag.setAttribute(o3, OperatorContext.STATS_LISTENERS,
      Lists.<StatsListener>newArrayList(mock(StatsListenerWithContext.class)));

  StatsListenerOperatorOld o4 = dag.addOperator("o4", new StatsListenerOperatorOld());

  PhysicalPlan plan = new PhysicalPlan(dag, new TestPlanContext());

  PTOperator p1 = plan.getOperators(dag.getMeta(o1)).get(0);
  StatsListener l = p1.statsListeners.get(0);
  Assert.assertTrue("Operator stats listener is wrapped ", l instanceof StatsListenerWithContext);

  PTOperator p2 = plan.getOperators(dag.getMeta(o2)).get(0);
  l = p1.statsListeners.get(0);
  Assert.assertTrue("Operator stats listener is wrapped ", l instanceof StatsListenerWithContext);

  PTOperator p3 = plan.getOperators(dag.getMeta(o3)).get(0);
  l = p1.statsListeners.get(0);
  Assert.assertTrue("Operator stats listener is wrapped ", l instanceof StatsListenerWithContext);

  PTOperator p4 = plan.getOperators(dag.getMeta(o4)).get(0);
  l = p1.statsListeners.get(0);
  Assert.assertTrue("Operator stats listener is wrapped ", l instanceof StatsListenerWithContext);
}
 

/**
 * Generate a conflict, Add a top level operator with name "m1_O1",
 * and add a module "m1" which will populate operator "O1", causing name conflict with
 * top level operator.
 */
@Test(expected = java.lang.IllegalArgumentException.class)
public void conflictingNamesWithExpandedModule()
{
  Configuration conf = new Configuration(false);
  LogicalPlanConfiguration lpc = new LogicalPlanConfiguration(conf);
  LogicalPlan dag = new LogicalPlan();
  DummyInputOperator in = dag.addOperator(componentName("m1", "O1"), new DummyInputOperator());
  Level2ModuleA module = dag.addModule("m1", new Level2ModuleA());
  dag.addStream("s1", in.out, module.mIn);
  lpc.prepareDAG(dag, null, "ModuleApp");
  dag.validate();
}
 

@Test
public void validatePositiveOiOOptionalInput()
{
  LogicalPlan plan = new LogicalPlan();
  RecoverableInputOperator inputOp1 = plan.addOperator("InputOperator1", new RecoverableInputOperator());
  GenericOperator genOp = plan.addOperator("GenericOperator", new GenericOperator());
  plan.addStream("OiO1", inputOp1.output, genOp.ip1).setLocality(Locality.THREAD_LOCAL);

  try {
    plan.validate();
    Assert.assertTrue("OiO validation", true);
  } catch (ConstraintViolationException ex) {
    Assert.fail("OiO Single Connected InputPort");
  }
}
 

@Test
public void testNodeLocal()
{

  LogicalPlan dag = new LogicalPlan();
  dag.getAttributes().put(com.datatorrent.api.Context.DAGContext.APPLICATION_PATH, new File("target", LocalityTest.class.getName()).getAbsolutePath());
  dag.setAttribute(OperatorContext.STORAGE_AGENT, new MemoryStorageAgent());

  GenericTestOperator o1 = dag.addOperator("o1", GenericTestOperator.class);

  GenericTestOperator partitioned = dag.addOperator("partitioned", GenericTestOperator.class);
  dag.getMeta(partitioned).getAttributes().put(OperatorContext.PARTITIONER, new StatelessPartitioner<GenericTestOperator>(2));

  GenericTestOperator partitionedParallel = dag.addOperator("partitionedParallel", GenericTestOperator.class);

  dag.addStream("o1_outport1", o1.outport1, partitioned.inport1).setLocality(null);

  dag.addStream("partitioned_outport1", partitioned.outport1, partitionedParallel.inport2).setLocality(Locality.NODE_LOCAL);
  dag.setInputPortAttribute(partitionedParallel.inport2, PortContext.PARTITION_PARALLEL, true);

  GenericTestOperator single = dag.addOperator("single", GenericTestOperator.class);
  dag.addStream("partitionedParallel_outport1", partitionedParallel.outport1, single.inport1);

  int maxContainers = 7;
  dag.setAttribute(LogicalPlan.CONTAINERS_MAX_COUNT, maxContainers);

  StreamingContainerManager scm = new StreamingContainerManager(dag);
  Assert.assertEquals("number required containers", 6, scm.containerStartRequests.size());

  ResourceRequestHandler rr = new ResourceRequestHandler();

  int containerMem = 2000;
  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()));

  Map<PTContainer, String> requestedHosts = Maps.newHashMap();
  for (ContainerStartRequest csr : scm.containerStartRequests) {
    String host = rr.getHost(csr, true);
    csr.container.host = host;
    // update the node report
    if (host != null) {
      requestedHosts.put(csr.container, host);
      nr = nodeReports.get(host);
      nr.getUsed().setMemory(nr.getUsed().getMemory() + containerMem);
    }
  }

  Assert.assertEquals("" + requestedHosts, nodeReports.keySet(), Sets.newHashSet(requestedHosts.values()));

  for (Map.Entry<PTContainer, String> e : requestedHosts.entrySet()) {
    for (PTOperator oper : e.getKey().getOperators()) {
      if (oper.getNodeLocalOperators().getOperatorSet().size() > 1) {
        String expHost = null;
        for (PTOperator nodeLocalOper : oper.getNodeLocalOperators().getOperatorSet()) {
          Assert.assertNotNull("host null " + nodeLocalOper.getContainer(), nodeLocalOper.getContainer().host);
          if (expHost == null) {
            expHost = nodeLocalOper.getContainer().host;
          } else {
            Assert.assertEquals("expected same host " + nodeLocalOper, expHost, nodeLocalOper.getContainer().host);
          }
        }
      }
    }
  }

}
 

/**
 * Test covering scenario when only new partitions are added during dynamic partitioning and there
 * are no changes to existing partitions and partition mapping
 */
@Test
public void testAugmentedDynamicPartitioning()
{
  LogicalPlan dag = new LogicalPlan();

  TestGeneratorInputOperator o1 = dag.addOperator("o1", TestGeneratorInputOperator.class);
  dag.setOperatorAttribute(o1, OperatorContext.PARTITIONER, new TestAugmentingPartitioner<TestGeneratorInputOperator>(3));
  dag.setOperatorAttribute(o1, OperatorContext.STATS_LISTENERS, Lists.newArrayList((StatsListener)new PartitioningTest.PartitionLoadWatch()));
  OperatorMeta o1Meta = dag.getMeta(o1);

  GenericTestOperator o2 = dag.addOperator("o2", GenericTestOperator.class);
  OperatorMeta o2Meta = dag.getMeta(o2);

  dag.addStream("o1.outport1", o1.outport, o2.inport1);

  int maxContainers = 10;
  dag.setAttribute(LogicalPlan.CONTAINERS_MAX_COUNT, maxContainers);

  TestPlanContext ctx = new TestPlanContext();
  dag.setAttribute(OperatorContext.STORAGE_AGENT, ctx);

  PhysicalPlan plan = new PhysicalPlan(dag, ctx);
  Assert.assertEquals("number of containers", 4, plan.getContainers().size());

  List<PTOperator> o1ops = plan.getOperators(o1Meta);
  Assert.assertEquals("number of o1 operators", 3, o1ops.size());

  List<PTOperator> o2ops = plan.getOperators(o2Meta);
  Assert.assertEquals("number of o2 operators", 1, o2ops.size());
  Set<PTOperator> expUndeploy = Sets.newLinkedHashSet();
  expUndeploy.addAll(plan.getOperators(o2Meta));
  expUndeploy.add(plan.getOperators(o2Meta).get(0).upstreamMerge.values().iterator().next());

  for (int i = 0; i < 2; ++i) {
    PartitioningTest.PartitionLoadWatch.put(o1ops.get(i), 1);
    plan.onStatusUpdate(o1ops.get(i));
  }

  ctx.backupRequests = 0;
  ctx.events.remove(0).run();

  Assert.assertEquals("number of containers", 6, plan.getContainers().size());

  Assert.assertEquals("undeployed opertors", expUndeploy, ctx.undeploy);
}
 

@SuppressWarnings("SleepWhileInLoop")
private void baseTestForQueueSize(int maxTuples, TestCollectorStatsListener statsListener, DAG.Locality locality) throws Exception
{
  LogicalPlan dag = new LogicalPlan();
  String workingDir = new File("target/baseTestForQueueSize").getAbsolutePath();
  dag.setAttribute(Context.OperatorContext.STORAGE_AGENT, new AsyncFSStorageAgent(workingDir, null));
  dag.getAttributes().put(LogicalPlan.STREAMING_WINDOW_SIZE_MILLIS, 200);
  TestOperator testOper = dag.addOperator("TestOperator", TestOperator.class);
  testOper.setMaxTuples(maxTuples);

  TestCollector collector = dag.addOperator("Collector", new TestCollector());
  if (statsListener != null) {
    dag.setOperatorAttribute(collector, OperatorContext.STATS_LISTENERS, Arrays.asList(new StatsListener[]{statsListener}));
  }

  dag.addStream("TestTuples", testOper.outport, collector.inport1).setLocality(locality);

  StramLocalCluster lc = new StramLocalCluster(dag);
  lc.runAsync();
  StreamingContainerManager dnmgr = lc.getStreamingContainerManager();
  Map<Integer, PTOperator> operatorMap = dnmgr.getPhysicalPlan().getAllOperators();
  for (PTOperator p : operatorMap.values()) {
    StramTestSupport.waitForActivation(lc, p);
  }

  long startTms = System.currentTimeMillis();
  if (statsListener != null) {
    while (statsListener.collectorOperatorStats.isEmpty() && (StramTestSupport.DEFAULT_TIMEOUT_MILLIS > System.currentTimeMillis() - startTms)) {
      Thread.sleep(300);
      LOG.debug("Waiting for stats");
    }
  } else {
    while (collector.collectorOperatorStats.isEmpty() && (StramTestSupport.DEFAULT_TIMEOUT_MILLIS > System.currentTimeMillis() - startTms)) {
      Thread.sleep(300);
      LOG.debug("Waiting for stats");
    }
  }
  if (statsListener != null) {
    statsListener.validateStats();
  } else {
    collector.validateStats();
  }
  lc.shutdown();
}
 

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

  dag.addStream("o1_outport1", o1.outport1, o2.inport1);// .setLocality(Locality.NODE_LOCAL);

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