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

@Test
public void validateOiOImplementation() throws Exception
{
  LogicalPlan lp = new LogicalPlan();
  String workingDir = new File("target/validateOiOImplementation").getAbsolutePath();
  lp.setAttribute(Context.OperatorContext.STORAGE_AGENT, new AsyncFSStorageAgent(workingDir, null));
  TestInputOperator io = lp.addOperator("Input Operator", new TestInputOperator());
  FirstGenericOperator go = lp.addOperator("First Generic Operator", new FirstGenericOperator());
  SecondGenericOperator out = lp.addOperator("Second Generic Operator", new SecondGenericOperator());

  /*
   * This tests make sure that even if the application_window_count is different the endWindow() is called for
   * end_stream
   */
  lp.getOperatorMeta("Second Generic Operator").getAttributes().put(Context.OperatorContext.APPLICATION_WINDOW_COUNT, 2);
  StreamMeta stream = lp.addStream("Stream", io.output, go.input);
  StreamMeta stream1 = lp.addStream("Stream1", go.output, out.input);

  stream1.setLocality(Locality.THREAD_LOCAL);
  lp.validate();
  StramLocalCluster slc = new StramLocalCluster(lp);
  slc.run();
  Assert.assertEquals("End Window Count", FirstGenericOperator.endwindowCount, SecondGenericOperator.endwindowCount);
}
 

@Test
public void testNumberOfUnifiersWithEvenPartitions()
{
  LogicalPlan dag = new LogicalPlan();
  dag.setAttribute(OperatorContext.STORAGE_AGENT, new StramTestSupport.MemoryStorageAgent());
  GenericTestOperator node1 = dag.addOperator("node1", GenericTestOperator.class);
  GenericTestOperator node2 = dag.addOperator("node2", GenericTestOperator.class);
  dag.addStream("node1.outport1", node1.outport1, node2.inport1);
  dag.setOperatorAttribute(node1, OperatorContext.PARTITIONER, new StatelessPartitioner<GenericTestOperator>(8));
  dag.setOutputPortAttribute(node1.outport1, PortContext.UNIFIER_LIMIT, 4);
  PhysicalPlan plan = new PhysicalPlan(dag, new TestPlanContext());
  List<PTContainer> containers = plan.getContainers();
  int unifierCount = 0;
  int totalOperators = 0;
  for (PTContainer container : containers) {
    List<PTOperator> operators = container.getOperators();
    for (PTOperator operator : operators) {
      totalOperators++;
      if (operator.isUnifier()) {
        unifierCount++;
      }
    }
  }
  Assert.assertEquals("Number of operators", 12, totalOperators);
  Assert.assertEquals("Number of unifiers", 3, unifierCount);
}
 

@Test
public void testNumberOfUnifiers()
{
  LogicalPlan dag = new LogicalPlan();
  dag.setAttribute(OperatorContext.STORAGE_AGENT, new StramTestSupport.MemoryStorageAgent());
  GenericTestOperator node1 = dag.addOperator("node1", GenericTestOperator.class);
  GenericTestOperator node2 = dag.addOperator("node2", GenericTestOperator.class);
  dag.addStream("node1.outport1", node1.outport1, node2.inport1);
  dag.setOperatorAttribute(node1, OperatorContext.PARTITIONER, new StatelessPartitioner<GenericTestOperator>(5));
  dag.setOutputPortAttribute(node1.outport1, PortContext.UNIFIER_LIMIT, 3);
  PhysicalPlan plan = new PhysicalPlan(dag, new TestPlanContext());
  List<PTContainer> containers = plan.getContainers();
  int unifierCount = 0;
  int totalOperators = 0;
  for (PTContainer container : containers) {
    List<PTOperator> operators = container.getOperators();
    for (PTOperator operator : operators) {
      totalOperators++;
      if (operator.isUnifier()) {
        unifierCount++;
      }
    }
  }
  Assert.assertEquals("Number of operators", 8, totalOperators);
  Assert.assertEquals("Number of unifiers", 2, unifierCount);
}
 

@Test
public void testInputOperatorRecovery() throws Exception
{
  RecoverableInputOperator.initGenTuples();
  CollectorOperator.collection.clear();
  int maxTuples = 30;
  LogicalPlan dag = new LogicalPlan();
  String workingDir = new File("target/testInputOperatorRecovery").getAbsolutePath();
  AsyncFSStorageAgent asyncFSStorageAgent = new AsyncFSStorageAgent(workingDir, null);
  asyncFSStorageAgent.setSyncCheckpoint(true);
  dag.setAttribute(Context.OperatorContext.STORAGE_AGENT, asyncFSStorageAgent);
  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);
  dag.addStream("connection", rip.output, cm.input);

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

  Assert.assertEquals("Generated Outputs", maxTuples, CollectorOperator.collection.size());
}
 

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

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

/**
 * 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 testContainersForSlidingWindow()
{
  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.APPLICATION_WINDOW_COUNT, 4);
  dag.setOperatorAttribute(o1, OperatorContext.SLIDE_BY_WINDOW_COUNT, 2);
  dag.getOperatorMeta("o1").getMeta(o1.outport1).getUnifierMeta().getAttributes().put(OperatorContext.MEMORY_MB, 2000);
  dag.getOperatorMeta("o1").getMeta(o1.outport2).getUnifierMeta().getAttributes().put(OperatorContext.MEMORY_MB, 4000);

  dag.addStream("o1.outport1", o1.outport1, o2.inport1);
  dag.addStream("o1.outport2", o1.outport2, o2.inport2);
  dag.addStream("o2.outport1", o2.outport1, o3.inport1);
  PhysicalPlan plan = new PhysicalPlan(dag, new TestPlanContext());
  Assert.assertEquals("number of containers", 5, plan.getContainers().size());
  boolean sawOutput1Slider = false;
  boolean sawOutput2Slider = false;
  for (PTContainer container : plan.getContainers()) {
    Assert.assertEquals("number of operators in each container is 1", container.operators.size(), 1);
    if (container.operators.get(0).isUnifier()) {
      String name = container.operators.get(0).getName();
      if (name.equals("o1.outport1#slider")) {
        sawOutput1Slider = true;
        Assert.assertEquals("container memory is 2512", container.getRequiredMemoryMB(), 2512);
      } else if (name.equals("o1.outport2#slider")) {
        sawOutput2Slider = true;
        Assert.assertEquals("container memory is 2512", container.getRequiredMemoryMB(), 4512);
      }
    }
  }
  Assert.assertEquals("Found output1 slider", true, sawOutput1Slider);
  Assert.assertEquals("Found output2 slider", true, sawOutput2Slider);
}
 

/**
 * 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 testOutofSequenceError() throws Exception
{
  logger.info("Testing Out of Sequence Error");
  LogicalPlan dag = new LogicalPlan();
  String workingDir = new File("target/testOutofSequenceError").getAbsolutePath();
  dag.setAttribute(Context.OperatorContext.STORAGE_AGENT, new AsyncFSStorageAgent(workingDir, null));
  RandomNumberGenerator rng = dag.addOperator("random", new RandomNumberGenerator());
  MyLogger ml = dag.addOperator("logger", new MyLogger());

  dag.addStream("stream", rng.output, ml.input);

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

private void test(int applicationWindowCount, int slideByWindowCount) throws Exception
{
  LogicalPlan dag = new LogicalPlan();
  String workingDir = new File("target/sliderTest").getAbsolutePath();
  dag.setAttribute(Context.OperatorContext.STORAGE_AGENT, new AsyncFSStorageAgent(workingDir, null));
  dag.getAttributes().put(LogicalPlan.STREAMING_WINDOW_SIZE_MILLIS, 100);
  Input input = dag.addOperator("Input", new Input());
  Sum sum = dag.addOperator("Sum", new Sum());
  dag.setOperatorAttribute(sum, OperatorContext.APPLICATION_WINDOW_COUNT, applicationWindowCount);
  dag.setOperatorAttribute(sum, OperatorContext.SLIDE_BY_WINDOW_COUNT, slideByWindowCount);
  Validator validate = dag.addOperator("validator", new Validator());
  Validator.numbersValidated = 0;
  validate.numberOfIntegers = applicationWindowCount;
  validate.slideByNumbers = slideByWindowCount;
  dag.addStream("input-sum", input.defaultOutputPort, sum.inputPort);
  dag.addStream("sum-validator", sum.outputPort, validate.validate);
  StramLocalCluster lc = new StramLocalCluster(dag);
  lc.runAsync();

  long startTms = System.currentTimeMillis();
  while (StramTestSupport.DEFAULT_TIMEOUT_MILLIS > System.currentTimeMillis() - startTms) {
    if (validate.numbersValidated > 5) {
      break;
    }
    Thread.sleep(100);
  }
  lc.shutdown();
  Assert.assertTrue("numbers validated more than zero ", validate.numbersValidated > 0);
}
 

@Test
public void testRecordingFlow() throws Exception
{
  LogicalPlan dag = new LogicalPlan();
  dag.setAttribute(Context.OperatorContext.STORAGE_AGENT, new AsyncFSStorageAgent(testWorkDir.getAbsolutePath(), null));

  dag.getAttributes().put(LogicalPlan.APPLICATION_PATH, "file://" + testWorkDir.getAbsolutePath());
  dag.getAttributes().put(LogicalPlan.TUPLE_RECORDING_PART_FILE_SIZE, 1024);  // 1KB per part

  TestGeneratorInputOperator op1 = dag.addOperator("op1", TestGeneratorInputOperator.class);
  GenericTestOperator op2 = dag.addOperator("op2", GenericTestOperator.class);
  GenericTestOperator op3 = dag.addOperator("op3", GenericTestOperator.class);

  op1.setEmitInterval(100); // emit every 100 msec
  dag.addStream("stream1", op1.outport, op2.inport1);//.setInline(true);
  dag.addStream("stream2", op2.outport1, op3.inport1);//.setInline(true);

  final StramLocalCluster localCluster = new StramLocalCluster(dag);
  localCluster.runAsync();

  final PTOperator ptOp2 = localCluster.findByLogicalNode(dag.getMeta(op2));
  StramTestSupport.waitForActivation(localCluster, ptOp2);

  testRecordingOnOperator(localCluster, ptOp2);

  final PTOperator ptOp1 = localCluster.findByLogicalNode(dag.getMeta(op1));
  StramTestSupport.waitForActivation(localCluster, ptOp1);

  testRecordingOnOperator(localCluster, ptOp1);

  localCluster.shutdown();
}
 

@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
public void testStaticPartitioning()
{
  LogicalPlan dag = new LogicalPlan();
  dag.setAttribute(OperatorContext.STORAGE_AGENT, new StramTestSupport.MemoryStorageAgent());

  TestGeneratorInputOperator node0 = dag.addOperator("node0", TestGeneratorInputOperator.class);
  GenericTestOperator node1 = dag.addOperator("node1", GenericTestOperator.class);
  PartitioningTestOperator partitioned = dag.addOperator("partitioned", PartitioningTestOperator.class);
  partitioned.setPartitionCount(partitioned.partitionKeys.length);
  GenericTestOperator singleton1 = dag.addOperator("singleton1", GenericTestOperator.class);
  GenericTestOperator singleton2 = dag.addOperator("singleton2", GenericTestOperator.class);

  dag.addStream("n0.inport1", node0.outport, node1.inport1);
  dag.addStream("n1.outport1", node1.outport1, partitioned.inport1, partitioned.inportWithCodec);
  dag.addStream("mergeStream", partitioned.outport1, singleton1.inport1, singleton2.inport1);

  dag.setAttribute(LogicalPlan.CONTAINERS_MAX_COUNT, 2);

  OperatorMeta partitionedMeta = dag.getMeta(partitioned);

  dag.validate();

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

  Assert.assertEquals("number of containers", 2, plan.getContainers().size());
  Assert.assertNotNull("partition map", partitioned.partitions);
  Assert.assertEquals("partition map " + partitioned.partitions, 3, partitioned.partitions.size());

  List<PTOperator> n2Instances = plan.getOperators(partitionedMeta);
  Assert.assertEquals("partition instances " + n2Instances, partitioned.partitionKeys.length, n2Instances.size());
  for (int i = 0; i < n2Instances.size(); i++) {
    PTOperator po = n2Instances.get(i);
    Map<String, PTInput> inputsMap = new HashMap<>();
    for (PTInput input: po.getInputs()) {
      inputsMap.put(input.portName, input);
      Assert.assertEquals("partitions " + input, Sets.newHashSet(partitioned.partitionKeys[i]), input.partitions.partitions);
      //Assert.assertEquals("codec " + input.logicalStream, PartitioningTestStreamCodec.class, input.logicalStream.getCodecClass());
    }
    Assert.assertEquals("number inputs " + inputsMap, Sets.newHashSet(PartitioningTestOperator.IPORT1, PartitioningTestOperator.INPORT_WITH_CODEC), inputsMap.keySet());
  }

  Collection<PTOperator> unifiers = plan.getMergeOperators(partitionedMeta);
  Assert.assertEquals("number unifiers " + partitionedMeta, 0, unifiers.size());
}
 

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

  GenericTestOperator node1 = dag.addOperator("node1", GenericTestOperator.class);
  GenericTestOperator node2 = dag.addOperator("node2", GenericTestOperator.class);
  dag.addStream("node1.outport1", node1.outport1, node2.inport2, node2.inport1);

  int initialPartitionCount = 5;
  OperatorMeta node2Decl = dag.getMeta(node2);
  node2Decl.getAttributes().put(OperatorContext.PARTITIONER, new StatelessPartitioner<GenericTestOperator>(initialPartitionCount));

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

  List<PTOperator> n2Instances = plan.getOperators(node2Decl);
  Assert.assertEquals("partition instances " + n2Instances, initialPartitionCount, n2Instances.size());

  List<Integer> assignedPartitionKeys = Lists.newArrayList();

  for (int i = 0; i < n2Instances.size(); i++) {
    PTOperator n2Partition = n2Instances.get(i);
    Assert.assertNotNull("partition keys null: " + n2Partition, n2Partition.getPartitionKeys());
    Map<InputPort<?>, PartitionKeys> pkeys = n2Partition.getPartitionKeys();
    Assert.assertEquals("partition keys size: " + pkeys, 1, pkeys.size()); // one port partitioned
    InputPort<?> expectedPort = node2.inport2;
    Assert.assertEquals("partition port: " + pkeys, expectedPort, pkeys.keySet().iterator().next());

    Assert.assertEquals("partition mask: " + pkeys, "111", Integer.toBinaryString(pkeys.get(expectedPort).mask));
    Set<Integer> pks = pkeys.get(expectedPort).partitions;
    Assert.assertTrue("number partition keys: " + pkeys, pks.size() == 1 || pks.size() == 2);
    assignedPartitionKeys.addAll(pks);
  }

  int expectedMask = Integer.parseInt("111", 2);
  Assert.assertEquals("assigned partitions ", expectedMask + 1, assignedPartitionKeys.size());
  for (int i = 0; i <= expectedMask; i++) {
    Assert.assertTrue("" + assignedPartitionKeys, assignedPartitionKeys.contains(i));
  }

}
 

@Test
public void testPartitionLocality()
{
  int partitionCount = 3;
  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);

  GenericTestOperator partitioned = dag.addOperator("partitioned", GenericTestOperator.class);
  LocalityPartitioner partitioner = new LocalityPartitioner();
  partitioner.setPartitionCount(partitionCount);
  dag.getMeta(partitioned).getAttributes().put(OperatorContext.PARTITIONER, partitioner);
  dag.addStream("o1_outport1", o1.outport1, partitioned.inport1);

  StreamingContainerManager scm = new StreamingContainerManager(dag);

  ResourceRequestHandler rr = new ResourceRequestHandler();

  int containerMem = 1000;
  Map<String, NodeReport> nodeReports = Maps.newHashMap();
  for (int i = 0; i < partitionCount; i++) {
    NodeReport nr = BuilderUtils.newNodeReport(BuilderUtils.newNodeId("host" + (i + 1), 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()));
  Set<String> expectedHosts = Sets.newHashSet();
  for (int i = 0; i < partitionCount; i++) {
    expectedHosts.add("host" + (i + 1));
  }
  for (ContainerStartRequest csr : scm.containerStartRequests) {
    String host = rr.getHost(csr, true);
    if (host != null) {
      expectedHosts.remove(host);
    }
  }
  Assert.assertTrue("All the allocated hosts removed", expectedHosts.isEmpty());

}
 

@Test
public void testRepartitioningScaleUp()
{
  LogicalPlan dag = new LogicalPlan();

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

  dag.addStream("o1.outport1", o1.outport1, o2.inport1, o2.inport2);
  dag.addStream("mergeStream", o2.outport1, mergeNode.inport1);

  OperatorMeta o2Meta = dag.getMeta(o2);
  o2Meta.getAttributes().put(OperatorContext.STATS_LISTENERS,
      Lists.newArrayList((StatsListener)new PartitionLoadWatch(0, 5)));
  o2Meta.getAttributes().put(OperatorContext.PARTITIONER, new StatelessPartitioner<GenericTestOperator>(1));

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

  Assert.assertEquals("number of operators", 3, plan.getAllOperators().size());
  Assert.assertEquals("number of save requests", 3, ctx.backupRequests);

  List<PTOperator> o2Partitions = plan.getOperators(o2Meta);
  Assert.assertEquals("partition count " + o2Meta, 1, o2Partitions.size());

  PTOperator o2p1 = o2Partitions.get(0);
  Assert.assertEquals("stats handlers " + o2p1, 1, o2p1.statsListeners.size());
  StatsListener sl = o2p1.statsListeners.get(0);
  Assert.assertTrue("stats handlers " + o2p1.statsListeners, sl instanceof PartitionLoadWatch);
  ((PartitionLoadWatch)sl).evalIntervalMillis = -1; // no delay

  setThroughput(o2p1, 10);
  plan.onStatusUpdate(o2p1);
  Assert.assertEquals("partitioning triggered", 1, ctx.events.size());
  ctx.backupRequests = 0;
  ctx.events.remove(0).run();

  o2Partitions = plan.getOperators(o2Meta);
  Assert.assertEquals("partition count " + o2Partitions, 2, o2Partitions.size());
  o2p1 = o2Partitions.get(0);
  Assert.assertEquals("sinks " + o2p1.getOutputs(), 1, o2p1.getOutputs().size());
  PTOperator o2p2 = o2Partitions.get(1);
  Assert.assertEquals("sinks " + o2p2.getOutputs(), 1, o2p2.getOutputs().size());

  Set<PTOperator> expUndeploy = Sets.newHashSet(plan.getOperators(dag.getMeta(mergeNode)));
  expUndeploy.add(o2p1);

  expUndeploy.addAll(plan.getOperators(dag.getMeta(mergeNode)).get(0).upstreamMerge.values());

  // verify load update generates expected events per configuration

  setThroughput(o2p1, 0);
  plan.onStatusUpdate(o2p1);
  Assert.assertEquals("load min", 0, ctx.events.size());

  setThroughput(o2p1, 3);
  plan.onStatusUpdate(o2p1);
  Assert.assertEquals("load within range", 0, ctx.events.size());

  setThroughput(o2p1, 10);
  plan.onStatusUpdate(o2p1);
  Assert.assertEquals("load exceeds max", 1, ctx.events.size());

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

  Assert.assertEquals("new partitions", 3, plan.getOperators(o2Meta).size());
  Assert.assertTrue("", plan.getOperators(o2Meta).contains(o2p2));

  for (PTOperator partition : plan.getOperators(o2Meta)) {
    Assert.assertNotNull("container null " + partition, partition.getContainer());
    Assert.assertEquals("outputs " + partition, 1, partition.getOutputs().size());
    Assert.assertEquals("downstream operators " + partition.getOutputs().get(0).sinks, 1, partition.getOutputs().get(0).sinks.size());
  }
  Assert.assertEquals("" + ctx.undeploy, expUndeploy, ctx.undeploy);

  Set<PTOperator> expDeploy = Sets.newHashSet(plan.getOperators(dag.getMeta(mergeNode)));
  expDeploy.addAll(plan.getOperators(o2Meta));
  expDeploy.remove(o2p2);

  expDeploy.addAll(plan.getOperators(dag.getMeta(mergeNode)).get(0).upstreamMerge.values());

  Assert.assertEquals("" + ctx.deploy, expDeploy, ctx.deploy);
  Assert.assertEquals("Count of storage requests", 2, ctx.backupRequests);

  // partitioning skipped on insufficient head room
  o2p1 = plan.getOperators(o2Meta).get(0);
  plan.setAvailableResources(0);
  setThroughput(o2p1, 10);
  plan.onStatusUpdate(o2p1);
  Assert.assertEquals("not repartitioned", 1, ctx.events.size());
  ctx.events.remove(0).run();
  Assert.assertEquals("partition count unchanged", 3, plan.getOperators(o2Meta).size());

}
 

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

/**
 * Test unifier gets removed when number partitions drops to 1.
 */
@Test
public void testRepartitioningScaleDownSinglePartition()
{
  LogicalPlan dag = new LogicalPlan();

  TestInputOperator<?> o1 = dag.addOperator("o1", TestInputOperator.class);
  GenericTestOperator o2 = dag.addOperator("o2", GenericTestOperator.class);

  dag.addStream("o1.outport1", o1.output, o2.inport1);
  OperatorMeta o1Meta = dag.getMeta(o1);
  dag.setOperatorAttribute(o1, OperatorContext.PARTITIONER, new StatelessPartitioner<GenericTestOperator>(2));
  dag.setOperatorAttribute(o1, OperatorContext.STATS_LISTENERS, Arrays.asList(new StatsListener[]{new PartitioningTest.PartitionLoadWatch()}));

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

  List<PTOperator> o1Partitions = plan.getOperators(o1Meta);
  Assert.assertEquals("partitions " + o1Partitions, 2, o1Partitions.size());
  PTOperator o1p1 = o1Partitions.get(0);
  PTOperator p1Doper = o1p1.getOutputs().get(0).sinks.get(0).target;
  Assert.assertSame("", p1Doper.getOperatorMeta(), o1Meta.getMeta(o1.output).getUnifierMeta());
  Assert.assertTrue("unifier ", p1Doper.isUnifier());
  Assert.assertEquals("Unifiers " + o1Meta, 1, o1p1.getOutputs().get(0).sinks.size());

  Collection<PTOperator> o1Unifiers = new ArrayList<>(plan.getOperators(dag.getMeta(o2)).get(0).upstreamMerge.values());

  StatsListener l = o1p1.statsListeners.get(0);
  Assert.assertTrue("stats handlers " + o1p1.statsListeners, l instanceof PartitioningTest.PartitionLoadWatch);
  PartitioningTest.PartitionLoadWatch.put(o1p1, -1);
  PartitioningTest.PartitionLoadWatch.put(o1Partitions.get(1), -1);

  plan.onStatusUpdate(o1p1);
  plan.onStatusUpdate(o1Partitions.get(1));
  Assert.assertEquals("partition scaling triggered", 1, ctx.events.size());
  ctx.events.remove(0).run();

  List<PTOperator> o1NewPartitions = plan.getOperators(o1Meta);
  Assert.assertEquals("partitions " + o1NewPartitions, 1, o1NewPartitions.size());

  List<PTOperator> o1NewUnifiers = new ArrayList<>(plan.getOperators(dag.getMeta(o2)).get(0).upstreamMerge.values());

  Assert.assertEquals("unifiers " + o1Meta, 0, o1NewUnifiers.size());
  p1Doper = o1p1.getOutputs().get(0).sinks.get(0).target;
  Assert.assertTrue("", p1Doper.getOperatorMeta() == dag.getMeta(o2));
  Assert.assertFalse("unifier ", p1Doper.isUnifier());

  Assert.assertTrue("removed unifier from deployment " + ctx.undeploy, ctx.undeploy.containsAll(o1Unifiers));
  Assert.assertFalse("removed unifier from deployment " + ctx.deploy, ctx.deploy.containsAll(o1Unifiers));

  // scale up, ensure unifier is setup at activation checkpoint
  setActivationCheckpoint(o1NewPartitions.get(0), 3);
  PartitioningTest.PartitionLoadWatch.put(o1NewPartitions.get(0), 1);
  plan.onStatusUpdate(o1NewPartitions.get(0));
  Assert.assertEquals("partition scaling triggered", 1, ctx.events.size());
  ctx.events.remove(0).run();

  o1NewUnifiers.addAll(plan.getOperators(dag.getMeta(o2)).get(0).upstreamMerge.values());

  Assert.assertEquals("unifiers " + o1Meta, 1, o1NewUnifiers.size());
  Assert.assertEquals("unifier activation checkpoint " + o1Meta, 3, o1NewUnifiers.get(0).recoveryCheckpoint.windowId);
}
 

@Test
@SuppressWarnings("unchecked")
public void testInline()
{

  LogicalPlan dag = new LogicalPlan();

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

  PartitioningTestOperator partOperator = dag.addOperator("partNode", PartitioningTestOperator.class);
  partOperator.partitionKeys = new Integer[] {0,1};
  dag.getMeta(partOperator).getAttributes().put(OperatorContext.PARTITIONER, new StatelessPartitioner<GenericTestOperator>(partOperator.partitionKeys.length));

  dag.addStream("o1_outport1", o1.outport1, o2.inport1, o3.inport1, partOperator.inport1)
          .setLocality(null);

  // same container for o2 and o3
  dag.addStream("o2_outport1", o2.outport1, o3.inport2)
      .setLocality(Locality.CONTAINER_LOCAL);

  dag.addStream("o3_outport1", o3.outport1, partOperator.inport2);

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

  PhysicalPlan plan = new PhysicalPlan(dag, new TestPlanContext());
  Assert.assertEquals("number of containers", maxContainers, plan.getContainers().size());
  Assert.assertEquals("operators container 0", 1, plan.getContainers().get(0).getOperators().size());

  Assert.assertEquals("operators container 0", 1, plan.getContainers().get(0).getOperators().size());
  Set<OperatorMeta> c2ExpNodes = Sets.newHashSet(dag.getMeta(o2), dag.getMeta(o3));
  Set<OperatorMeta> c2ActNodes = new HashSet<>();
  PTContainer c2 = plan.getContainers().get(1);
  for (PTOperator pNode : c2.getOperators()) {
    c2ActNodes.add(pNode.getOperatorMeta());
  }
  Assert.assertEquals("operators " + c2, c2ExpNodes, c2ActNodes);

  // one container per partition
  OperatorMeta partOperMeta = dag.getMeta(partOperator);
  List<PTOperator> partitions = plan.getOperators(partOperMeta);
  for (PTOperator partition : partitions) {
    Assert.assertEquals("operators container" + partition, 1, partition.getContainer().getOperators().size());
  }

}