com.datatorrent.stram.plan.logical.LogicalPlan Java Examples

private FinalVars(LogicalPlan dag, long tms)
{
  Attribute.AttributeMap attributes = dag.getAttributes();
  /* try to align to it to please eyes. */
  windowStartMillis = tms - (tms % 1000);

  if (attributes.get(LogicalPlan.APPLICATION_PATH) == null) {
    throw new IllegalArgumentException("Not set: " + LogicalPlan.APPLICATION_PATH);
  }

  this.appPath = attributes.get(LogicalPlan.APPLICATION_PATH);

  if (attributes.get(LogicalPlan.STREAMING_WINDOW_SIZE_MILLIS) == null) {
    attributes.put(LogicalPlan.STREAMING_WINDOW_SIZE_MILLIS, 500);
  }
  if (attributes.get(LogicalPlan.CHECKPOINT_WINDOW_COUNT) == null) {
    attributes.put(LogicalPlan.CHECKPOINT_WINDOW_COUNT, 30000 / attributes.get(LogicalPlan.STREAMING_WINDOW_SIZE_MILLIS));
  }

  this.heartbeatTimeoutMillis = dag.getValue(LogicalPlan.HEARTBEAT_TIMEOUT_MILLIS);
  this.maxWindowsBehindForStats = dag.getValue(LogicalPlan.STATS_MAX_ALLOWABLE_WINDOWS_LAG);
  this.enableStatsRecording = dag.getValue(LogicalPlan.ENABLE_STATS_RECORDING);
  this.rpcLatencyCompensationSamples = dag.getValue(LogicalPlan.RPC_LATENCY_COMPENSATION_SAMPLES);
}
 

/**
 * Return information about operators and inner modules of a module.
 *
 * @param dag        top level DAG
 * @param moduleMeta module information. DAG within module is used for constructing response.
 * @return
 */
private static Map<String, Object> getLogicalModuleDetails(LogicalPlan dag, LogicalPlan.ModuleMeta moduleMeta)
{
  Map<String, Object> moduleDetailMap = new HashMap<>();
  ArrayList<String> operatorArray = new ArrayList<>();
  moduleDetailMap.put("name", moduleMeta.getName());
  moduleDetailMap.put("className", moduleMeta.getGenericOperator().getClass().getName());

  moduleDetailMap.put("operators", operatorArray);
  for (OperatorMeta operatorMeta : moduleMeta.getDag().getAllOperators()) {
    if (operatorMeta.getModuleName() == null) {
      String fullName = moduleMeta.getFullName() + LogicalPlan.MODULE_NAMESPACE_SEPARATOR + operatorMeta.getName();
      operatorArray.add(fullName);
    }
  }
  ArrayList<Map<String, Object>> modulesArray = new ArrayList<>();
  moduleDetailMap.put("modules", modulesArray);
  for (LogicalPlan.ModuleMeta meta : moduleMeta.getDag().getAllModules()) {
    modulesArray.add(getLogicalModuleDetails(dag, meta));
  }
  return moduleDetailMap;
}
 

private void setInput(PTOperator oper, InputPortMeta ipm, PTOperator sourceOper, PartitionKeys pks)
{
  // TODO: see if this can be handled more efficiently
  for (PTInput in : oper.inputs) {
    if (in.source.source == sourceOper && in.logicalStream == streamMeta && ipm.getPortName().equals(in.portName)) {
      return;
    }
  }
  // link to upstream output(s) for this stream
  for (PTOutput upstreamOut : sourceOper.outputs) {
    if (upstreamOut.logicalStream == streamMeta) {
      PTInput input = new PTInput(ipm.getPortName(), streamMeta, oper, pks, upstreamOut, ipm.getValue(LogicalPlan.IS_CONNECTED_TO_DELAY_OPERATOR));
      oper.inputs.add(input);
    }
  }
}
 

@Test
public void testLoadFromPropertiesFile() throws IOException
{
  Properties props = new Properties();
  String resourcePath = "/testModuleTopology.properties";
  InputStream is = this.getClass().getResourceAsStream(resourcePath);
  if (is == null) {
    throw new RuntimeException("Could not load " + resourcePath);
  }
  props.load(is);
  LogicalPlanConfiguration pb = new LogicalPlanConfiguration(new Configuration(false))
      .addFromProperties(props, null);

  LogicalPlan dag = new LogicalPlan();
  pb.populateDAG(dag);
  pb.prepareDAG(dag, null, "testApplication");
  dag.validate();
  validateTopLevelOperators(dag);
  validateTopLevelStreams(dag);
  validatePublicMethods(dag);
}
 

@Test
public void testSQLSelectInsertWithAPI() throws IOException, ClassNotFoundException
{
  LogicalPlan dag = new LogicalPlan();

  String schemaIn = "{\"separator\":\",\",\"quoteChar\":\"\\\"\",\"fields\":[" +
      "{\"name\":\"RowTime\",\"type\":\"Date\",\"constraints\":{\"format\":\"dd/MM/yyyy hh:mm:ss\"}}," +
      "{\"name\":\"id\",\"type\":\"Integer\"}," +
      "{\"name\":\"Product\",\"type\":\"String\"}," +
      "{\"name\":\"units\",\"type\":\"Integer\"}]}";
  String schemaOut = "{\"separator\":\",\",\"quoteChar\":\"\\\"\",\"fields\":[" +
      "{\"name\":\"RowTime\",\"type\":\"Date\",\"constraints\":{\"format\":\"dd/MM/yyyy hh:mm:ss\"}}," +
      "{\"name\":\"Product\",\"type\":\"String\"}]}";

  SQLExecEnvironment.getEnvironment()
      .registerTable("ORDERS", new FileEndpoint("dummyFilePathInput", new CSVMessageFormat(schemaIn)))
      .registerTable("SALES", new FileEndpoint("dummyFilePathOutput", "out.tmp", new CSVMessageFormat(schemaOut)))
      .executeSQL(dag, "INSERT INTO SALES SELECT STREAM FLOOR(ROWTIME TO HOUR), SUBSTRING(PRODUCT, 0, 5) " +
      "FROM ORDERS WHERE id > 3");

  dag.validate();
}
 

private void setTokenRefreshCredentials(LogicalPlan dag, Configuration conf) throws IOException
{
  String principal = conf.get(StramClientUtils.TOKEN_REFRESH_PRINCIPAL, StramUserLogin.getPrincipal());
  String keytabPath = conf.get(StramClientUtils.TOKEN_REFRESH_KEYTAB, conf.get(StramClientUtils.KEY_TAB_FILE));
  if (keytabPath == null) {
    String keytab = StramUserLogin.getKeytab();
    if (keytab != null) {
      Path localKeyTabPath = new Path(keytab);
      try (FileSystem fs = StramClientUtils.newFileSystemInstance(conf)) {
        Path destPath = new Path(StramClientUtils.getApexDFSRootDir(fs, conf), localKeyTabPath.getName());
        if (!fs.exists(destPath)) {
          fs.copyFromLocalFile(false, false, localKeyTabPath, destPath);
        }
        keytabPath = destPath.toString();
      }
    }
  }
  LOG.debug("User principal is {}, keytab is {}", principal, keytabPath);
  if ((principal != null) && (keytabPath != null)) {
    dag.setAttribute(LogicalPlan.PRINCIPAL, principal);
    dag.setAttribute(LogicalPlan.KEY_TAB_FILE, keytabPath);
  } else {
    LOG.warn("Credentials for refreshing tokens not available, application may not be able to run indefinitely");
  }
}
 

/**
 * Remove the given partition with any associated parallel partitions and
 * per-partition outputStreams.
 *
 * @param oper
 * @return
 */
private void removePartition(PTOperator oper, PMapping operatorMapping)
{
  // remove any parallel partition
  for (PTOutput out : oper.outputs) {
    // copy list as it is modified by recursive remove
    for (PTInput in : Lists.newArrayList(out.sinks)) {
      for (LogicalPlan.InputPortMeta im : in.logicalStream.getSinks()) {
        PMapping m = this.logicalToPTOperator.get(im.getOperatorMeta());
        if (m.parallelPartitions == operatorMapping.parallelPartitions) {
          // associated operator parallel partitioned
          removePartition(in.target, operatorMapping);
          m.partitions.remove(in.target);
        }
      }
    }
  }
  // remove the operator
  removePTOperator(oper);
}
 

@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 testAddAttributeToArgs() throws Exception
{
  LogicalPlan dag = new LogicalPlan();
  dag.setAttribute(LogicalPlan.APPLICATION_NAME, APP_NAME);
  AddAttributeToArgsOperator operator = dag.addOperator("test", AddAttributeToArgsOperator.class);
  dag.getContextAttributes(operator).put(OperatorContext.RECOVERY_ATTEMPTS, 0);

  StramClient client = new StramClient(conf, dag);
  try {
    client.start();
    client.startApplication();
    Assert.assertTrue(client.monitorApplication());
  } finally {
    client.stop();
  }
}
 

private FinalVars(LogicalPlan dag, long tms)
{
  Attribute.AttributeMap attributes = dag.getAttributes();
  /* try to align to it to please eyes. */
  windowStartMillis = tms - (tms % 1000);

  if (attributes.get(LogicalPlan.APPLICATION_PATH) == null) {
    throw new IllegalArgumentException("Not set: " + LogicalPlan.APPLICATION_PATH);
  }

  this.appPath = attributes.get(LogicalPlan.APPLICATION_PATH);

  if (attributes.get(LogicalPlan.STREAMING_WINDOW_SIZE_MILLIS) == null) {
    attributes.put(LogicalPlan.STREAMING_WINDOW_SIZE_MILLIS, 500);
  }
  if (attributes.get(LogicalPlan.CHECKPOINT_WINDOW_COUNT) == null) {
    attributes.put(LogicalPlan.CHECKPOINT_WINDOW_COUNT, 30000 / attributes.get(LogicalPlan.STREAMING_WINDOW_SIZE_MILLIS));
  }

  this.heartbeatTimeoutMillis = dag.getValue(LogicalPlan.HEARTBEAT_TIMEOUT_MILLIS);
  this.maxWindowsBehindForStats = dag.getValue(LogicalPlan.STATS_MAX_ALLOWABLE_WINDOWS_LAG);
  this.enableStatsRecording = dag.getValue(LogicalPlan.ENABLE_STATS_RECORDING);
  this.rpcLatencyCompensationSamples = dag.getValue(LogicalPlan.RPC_LATENCY_COMPENSATION_SAMPLES);
}
 

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

@Override
public void setup(Context ctx)
{
  tupleRecordingPartFileSize = ctx.getValue(LogicalPlan.TUPLE_RECORDING_PART_FILE_SIZE);
  tupleRecordingPartFileTimeMillis = ctx.getValue(LogicalPlan.TUPLE_RECORDING_PART_FILE_TIME_MILLIS);
  appId = ctx.getValue(LogicalPlan.APPLICATION_ID);
  appPath = ctx.getValue(LogicalPlan.APPLICATION_PATH);
  codecs = ctx.getAttributes().get(Context.DAGContext.STRING_CODECS);

  wsClient = new PubSubWebSocketClientBuilder().setContext(ctx).build();

  RequestDelegateImpl impl = new RequestDelegateImpl();
  RequestFactory rf = ctx.getValue(ContainerContext.REQUEST_FACTORY);
  if (rf == null) {
    logger.warn("No request factory defined, recording is disabled!");
  } else {
    rf.registerDelegate(StramToNodeRequest.RequestType.START_RECORDING, impl);
    rf.registerDelegate(StramToNodeRequest.RequestType.STOP_RECORDING, impl);
    rf.registerDelegate(StramToNodeRequest.RequestType.SYNC_RECORDING, impl);
  }
}
 

private Configuration readLaunchConfiguration() throws IOException
{
  Path appPath = new Path(appContext.getApplicationPath());
  Path  configFilePath = new Path(appPath, LogicalPlan.LAUNCH_CONFIG_FILE_NAME);
  try {
    LOG.debug("Reading launch configuration file ");
    URI uri = appPath.toUri();
    Configuration config = new YarnConfiguration();
    fileContext = uri.getScheme() == null ? FileContext.getFileContext(config) : FileContext.getFileContext(uri, config);
    FSDataInputStream is = fileContext.open(configFilePath);
    config.addResource(is);
    LOG.debug("Read launch configuration");
    return config;
  } catch (FileNotFoundException ex) {
    LOG.warn("Configuration file not found {}", configFilePath);
    return new Configuration();
  }
}
 

@Override
public void setup(Context ctx)
{
  tupleRecordingPartFileSize = ctx.getValue(LogicalPlan.TUPLE_RECORDING_PART_FILE_SIZE);
  tupleRecordingPartFileTimeMillis = ctx.getValue(LogicalPlan.TUPLE_RECORDING_PART_FILE_TIME_MILLIS);
  appId = ctx.getValue(LogicalPlan.APPLICATION_ID);
  appPath = ctx.getValue(LogicalPlan.APPLICATION_PATH);
  codecs = ctx.getAttributes().get(Context.DAGContext.STRING_CODECS);

  wsClient = new PubSubWebSocketClientBuilder().setContext(ctx).build();

  RequestDelegateImpl impl = new RequestDelegateImpl();
  RequestFactory rf = ctx.getValue(ContainerContext.REQUEST_FACTORY);
  if (rf == null) {
    logger.warn("No request factory defined, recording is disabled!");
  } else {
    rf.registerDelegate(StramToNodeRequest.RequestType.START_RECORDING, impl);
    rf.registerDelegate(StramToNodeRequest.RequestType.STOP_RECORDING, impl);
    rf.registerDelegate(StramToNodeRequest.RequestType.SYNC_RECORDING, impl);
  }
}
 

private void testDownStreamPartition(Locality locality) throws Exception
{
  TestGeneratorInputOperator o1 = dag.addOperator("o1", TestGeneratorInputOperator.class);
  GenericTestOperator o2 = dag.addOperator("o2", GenericTestOperator.class);
  dag.setOperatorAttribute(o2, OperatorContext.PARTITIONER, new StatelessPartitioner<GenericTestOperator>(2));
  dag.addStream("o1Output1", o1.outport, o2.inport1).setLocality(locality);

  int maxContainers = 5;
  dag.setAttribute(LogicalPlan.CONTAINERS_MAX_COUNT, maxContainers);
  dag.setAttribute(OperatorContext.STORAGE_AGENT, new StramTestSupport.MemoryStorageAgent());
  dag.validate();
  PhysicalPlan plan = new PhysicalPlan(dag, new TestPlanContext());
  Assert.assertEquals("number of containers", 1, plan.getContainers().size());

  PTContainer container1 = plan.getContainers().get(0);
  Assert.assertEquals("number operators " + container1, 3, container1.getOperators().size());
  StramLocalCluster slc = new StramLocalCluster(dag);
  slc.run(5000);
}
 

public void testLinearInlineOperatorsRecovery() throws Exception
{
  RecoverableInputOperator.initGenTuples();
  CollectorOperator.collection.clear();
  CollectorOperator.duplicates.clear();

  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.getMeta(cm).getAttributes().put(OperatorContext.PROCESSING_MODE, processingMode);

  dag.addStream("connection", rip.output, cm.input).setLocality(Locality.CONTAINER_LOCAL);

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

private static List<PTOperator> assertNumberPartitions(final int count, final StramLocalCluster lc, final LogicalPlan.OperatorMeta ow) throws Exception
{
  WaitCondition c = new WaitCondition()
  {
    @Override
    public boolean isComplete()
    {
      List<PTOperator> operators = lc.getPlanOperators(ow);
      LOG.debug("Number of operators {}, expected number {}", operators.size(), count);
      return (operators.size() == count);
    }

  };
  StramTestSupport.awaitCompletion(c, 10000);
  Assert.assertTrue("Number partitions match " + ow, c.isComplete());
  return lc.getPlanOperators(ow);
}
 

public static InputPortMeta getInputPortMeta(LogicalPlan.OperatorMeta operatorMeta, StreamMeta streamMeta)
{
  InputPortMeta inputPortMeta = null;
  Map<InputPortMeta, StreamMeta> inputStreams = operatorMeta.getInputStreams();
  for (Map.Entry<InputPortMeta, StreamMeta> entry : inputStreams.entrySet()) {
    if (entry.getValue() == streamMeta) {
      inputPortMeta = entry.getKey();
      break;
    }
  }
  return inputPortMeta;
}
 

private void validatePublicMethods(LogicalPlan dag)
{
  // Logical dag contains 4 modules added on top level.
  List<String> moduleNames = new ArrayList<>();
  for (LogicalPlan.ModuleMeta moduleMeta : dag.getAllModules()) {
    moduleNames.add(moduleMeta.getName());
  }
  Assert.assertTrue(moduleNames.contains("Ma"));
  Assert.assertTrue(moduleNames.contains("Mb"));
  Assert.assertTrue(moduleNames.contains("Mc"));
  Assert.assertTrue(moduleNames.contains("Md"));
  Assert.assertTrue(moduleNames.contains("Me"));
  Assert.assertEquals("Number of modules are 5", 5, dag.getAllModules().size());

}
 

public void setAntiAffinityForContainers(LogicalPlan dag, Collection<AffinityRule> affinityRules, Map<PTOperator, PTContainer> operatorContainerMap)
{
  for (AffinityRule rule : affinityRules) {
    if (rule.getOperatorsList() != null && rule.getType() == Type.ANTI_AFFINITY) {
      for (int i = 0; i < rule.getOperatorsList().size() - 1; i++) {
        for (int j = i + 1; j < rule.getOperatorsList().size(); j++) {
          OperatorPair operators = new OperatorPair(rule.getOperatorsList().get(i), rule.getOperatorsList().get(j));

          PMapping firstPMapping = logicalToPTOperator.get(dag.getOperatorMeta(operators.first));
          OperatorMeta opMeta = dag.getOperatorMeta(operators.second);
          PMapping secondMapping = logicalToPTOperator.get(opMeta);
          for (PTOperator firstPtOperator : firstPMapping.partitions) {
            PTContainer firstContainer = operatorContainerMap.get(firstPtOperator);
            for (PTOperator secondPtOperator : secondMapping.partitions) {
              PTContainer secondContainer = operatorContainerMap.get(secondPtOperator);
              if (firstContainer == secondContainer || firstContainer.getStrictAntiPrefs().contains(secondContainer)) {
                continue;
              }
              if (rule.isRelaxLocality()) {
                firstContainer.getPreferredAntiPrefs().add(secondContainer);
                secondContainer.getPreferredAntiPrefs().add(firstContainer);
              } else {
                firstContainer.getStrictAntiPrefs().add(secondContainer);
                secondContainer.getStrictAntiPrefs().add(firstContainer);
              }
            }
          }
        }
      }
    }
  }
}
 

@Test
public void testJoin() throws IOException, ClassNotFoundException
{
  LogicalPlan dag = new LogicalPlan();
  String schemaIn0 = "{\"separator\":\",\",\"quoteChar\":\"\\\"\",\"fields\":[" +
      "{\"name\":\"RowTime\",\"type\":\"Date\",\"constraints\":{\"format\":\"dd/MM/yyyy hh:mm:ss Z\"}}," +
      "{\"name\":\"id\",\"type\":\"Integer\"}," +
      "{\"name\":\"Product\",\"type\":\"String\"}," +
      "{\"name\":\"units\",\"type\":\"Integer\"}]}";
  String schemaIn1 = "{\"separator\":\",\",\"quoteChar\":\"\\\"\",\"fields\":[" +
      "{\"name\":\"id\",\"type\":\"Integer\"}," +
      "{\"name\":\"Category\",\"type\":\"String\"}]}";
  String schemaOut = "{\"separator\":\",\",\"quoteChar\":\"\\\"\",\"fields\":[" +
      "{\"name\":\"RowTime1\",\"type\":\"Date\",\"constraints\":{\"format\":\"dd/MM/yyyy hh:mm:ss Z\"}}," +
      "{\"name\":\"RowTime2\",\"type\":\"Date\",\"constraints\":{\"format\":\"dd/MM/yyyy hh:mm:ss Z\"}}," +
      "{\"name\":\"Product\",\"type\":\"String\"}," +
      "{\"name\":\"Category\",\"type\":\"String\"}]}";

  String sql = "INSERT INTO SALES SELECT STREAM A.ROWTIME, FLOOR(A.ROWTIME TO DAY), " +
      "APEXCONCAT('OILPAINT', SUBSTRING(A.PRODUCT, 6, 7)), B.CATEGORY " +
      "FROM ORDERS AS A " +
      "JOIN CATEGORY AS B ON A.id = B.id " +
      "WHERE A.id > 3 AND A.PRODUCT LIKE 'paint%'";

  SQLExecEnvironment.getEnvironment()
      .registerTable("ORDERS", new KafkaEndpoint("localhost:9092", "testdata0", new CSVMessageFormat(schemaIn0)))
      .registerTable("CATEGORY", new KafkaEndpoint("localhost:9092", "testdata1", new CSVMessageFormat(schemaIn1)))
      .registerTable("SALES", new KafkaEndpoint("localhost:9092", "testresult", new CSVMessageFormat(schemaOut)))
      .registerFunction("APEXCONCAT", FileEndpointTest.class, "apex_concat_str")
      .executeSQL(dag, sql);

  dag.validate();
}
 

@Override
public LogicalPlan createApp(LogicalPlanConfiguration conf)
{
  try {
    return conf.createFromJson(json, getName());
  } catch (Exception e) {
    throw new IllegalArgumentException("Failed to load: " + this + "\n" + e.getMessage(), e);
  }
}
 

public StreamingContainerManager(LogicalPlan dag, boolean enableEventRecording, Clock clock)
{
  this.clock = clock;
  this.vars = new FinalVars(dag, clock.getTime());
  poolExecutor = Executors.newFixedThreadPool(4);
  // setup prior to plan creation for event recording
  if (enableEventRecording) {
    this.eventBus = new MBassador<>(BusConfiguration.Default(1, 1, 1));
  }
  this.plan = new PhysicalPlan(dag, this);
  this.journal = new Journal(this);
  init(enableEventRecording);
}
 

/**
 * For dry run on logical plan only
 * @param logicalPlan
 */
public PlanModifier(LogicalPlan logicalPlan)
{
  this.logicalPlan = logicalPlan;
  this.physicalPlan = null;
  init();
}
 

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

public StreamingContainerManager(LogicalPlan dag, boolean enableEventRecording, Clock clock)
{
  this.clock = clock;
  this.vars = new FinalVars(dag, clock.getTime());
  poolExecutor = Executors.newFixedThreadPool(4);
  // setup prior to plan creation for event recording
  if (enableEventRecording) {
    this.eventBus = new MBassador<>(BusConfiguration.Default(1, 1, 1));
  }
  this.plan = new PhysicalPlan(dag, this);
  this.journal = new Journal(this);
  init(enableEventRecording);
}
 

private FinalVars(FinalVars other, LogicalPlan dag)
{
  this.windowStartMillis = other.windowStartMillis;
  this.heartbeatTimeoutMillis = other.heartbeatTimeoutMillis;
  this.maxWindowsBehindForStats = other.maxWindowsBehindForStats;
  this.enableStatsRecording = other.enableStatsRecording;
  this.appPath = dag.getValue(LogicalPlan.APPLICATION_PATH);
  this.rpcLatencyCompensationSamples = other.rpcLatencyCompensationSamples;
}
 

private void setupRecording(boolean enableEventRecording)
{
  if (this.vars.enableStatsRecording) {
    statsRecorder = new FSStatsRecorder();
    statsRecorder.setBasePath(this.vars.appPath + "/" + LogicalPlan.SUBDIR_STATS);
    statsRecorder.setup();
  }
  if (enableEventRecording) {
    eventRecorder = new FSEventRecorder(plan.getLogicalPlan().getValue(LogicalPlan.APPLICATION_ID));
    eventRecorder.setBasePath(this.vars.appPath + "/" + LogicalPlan.SUBDIR_EVENTS);
    eventRecorder.setWebSocketClient(wsClient);
    eventRecorder.setup();
    eventBus.subscribe(eventRecorder);
  }
}
 

@Test
public void testContainerCores()
{
  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);
  GenericTestOperator o4 = dag.addOperator("o4", GenericTestOperator.class);
  GenericTestOperator o5 = dag.addOperator("o5", GenericTestOperator.class);
  GenericTestOperator o6 = dag.addOperator("o6", GenericTestOperator.class);
  dag.setOperatorAttribute(o1,OperatorContext.VCORES,1);
  dag.setOperatorAttribute(o2,OperatorContext.VCORES,2);
  dag.setOperatorAttribute(o3,OperatorContext.VCORES,3);
  dag.setOperatorAttribute(o4,OperatorContext.VCORES,4);
  dag.setOperatorAttribute(o5,OperatorContext.VCORES,5);
  dag.setOperatorAttribute(o6,OperatorContext.VCORES,6);

  dag.addStream("o1.outport1", o1.outport1, o2.inport1).setLocality(Locality.CONTAINER_LOCAL);
  dag.addStream("o2.outport1", o2.outport1, o3.inport1, o4.inport1).setLocality(Locality.THREAD_LOCAL);
  dag.addStream("o3.output1", o3.outport1, o5.inport1).setLocality(Locality.THREAD_LOCAL);
  dag.addStream("o4.output1", o4.outport1, o5.inport2).setLocality(Locality.THREAD_LOCAL);
  dag.addStream("o5.output1", o5.outport1, o6.inport1).setLocality(Locality.CONTAINER_LOCAL);

  dag.setAttribute(LogicalPlan.CONTAINERS_MAX_COUNT, 2);

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

  Assert.assertEquals("number of containers", 1, plan.getContainers().size());
  Assert.assertEquals("vcores container 1 is 12", 12, plan.getContainers().get(0).getRequiredVCores());
}