Java Code Examples for org.apache.samza.runtime.LocalApplicationRunner#waitForFinish()

/**
 * Run the application with the specified timeout
 *
 * @param timeout time to wait for the application to finish. This timeout does not include
 *                input stream initialization time or the assertion time over output streams. This timeout just accounts
 *                for time that samza job takes run. Timeout must be greater than 0.
 * @throws SamzaException if Samza job fails with exception and returns UnsuccessfulFinish as the statuscode
 */
public void run(Duration timeout) {
  Preconditions.checkNotNull(app);
  Preconditions.checkState(!timeout.isZero() || !timeout.isNegative(), "Timeouts should be positive");
  // Cleaning store directories to ensure current run does not pick up state from previous run
  deleteStoreDirectories();
  Config config = new MapConfig(JobPlanner.generateSingleJobConfig(configs));
  final LocalApplicationRunner runner = new LocalApplicationRunner(app, config, new InMemoryMetadataStoreFactory());
  runner.run(externalContext);
  if (!runner.waitForFinish(timeout)) {
    throw new SamzaException("Timed out waiting for application to finish");
  }
  ApplicationStatus status = runner.status();
  deleteStoreDirectories();
  if (status.getStatusCode() == ApplicationStatus.StatusCode.UnsuccessfulFinish) {
    throw new SamzaException("Application could not finish successfully", status.getThrowable());
  }
}
 

public static void main(String[] args) {
    String[] configArgs = {"--config-factory=org.apache.samza.config.factories.PropertiesConfigFactory"
            , "--config-path=samza-connector/src/main/Properties/config.properties"};
    CommandLine cmdLine = new CommandLine();
    OptionSet options = cmdLine.parser().parse(configArgs);
    Config config = cmdLine.loadConfig(options);
    LocalApplicationRunner runner = new LocalApplicationRunner(new SamzaSumDemo(), config);
    runner.run();
    runner.waitForFinish();
}
 

@Test
public void testTableWithLowLevelApi() throws Exception {
  Map<String, String> configs = getBaseJobConfig(bootstrapUrl(), zkConnect());
  configs.put("streams.PageView.samza.system", "test");
  configs.put("streams.PageView.source", Base64Serializer.serialize(TestTableData.generatePageViews(10)));
  configs.put("streams.PageView.partitionCount", String.valueOf(4));
  configs.put("task.inputs", "test.PageView");

  Config config = new MapConfig(configs);
  final LocalApplicationRunner runner = new LocalApplicationRunner(new MyTaskApplication(), config);
  executeRun(runner, config);
  runner.waitForFinish();
}
 

@Test
public void testSendTo() throws Exception {

  int count = 10;
  Profile[] profiles = TestTableData.generateProfiles(count);

  int partitionCount = 4;
  Map<String, String> configs = getBaseJobConfig(bootstrapUrl(), zkConnect());

  configs.put("streams.Profile.samza.system", "test");
  configs.put("streams.Profile.source", Base64Serializer.serialize(profiles));
  configs.put("streams.Profile.partitionCount", String.valueOf(partitionCount));

  MyMapFunction mapFn = new MyMapFunction();

  final StreamApplication app = appDesc -> {

    Table<KV<Integer, Profile>> table = appDesc.getTable(new InMemoryTableDescriptor("t1",
        KVSerde.of(new IntegerSerde(), new ProfileJsonSerde())));
    DelegatingSystemDescriptor ksd = new DelegatingSystemDescriptor("test");
    GenericInputDescriptor<Profile> isd = ksd.getInputDescriptor("Profile", new NoOpSerde<>());

    appDesc.getInputStream(isd)
        .map(mapFn)
        .sendTo(table);
  };

  Config config = new MapConfig(configs);
  final LocalApplicationRunner runner = new LocalApplicationRunner(app, config);
  executeRun(runner, config);
  runner.waitForFinish();

  for (int i = 0; i < partitionCount; i++) {
    MyMapFunction mapFnCopy = MyMapFunction.getMapFunctionByTask(String.format("Partition %d", i));
    assertEquals(count, mapFnCopy.received.size());
    mapFnCopy.received.forEach(p -> Assert.assertTrue(mapFnCopy.table.get(p.getMemberId()) != null));
  }
}
 

@Test
public void testStreamTableJoin() throws Exception {

  int count = 10;
  PageView[] pageViews = TestTableData.generatePageViews(count);
  Profile[] profiles = TestTableData.generateProfiles(count);

  int partitionCount = 4;
  Map<String, String> configs = getBaseJobConfig(bootstrapUrl(), zkConnect());

  configs.put("streams.PageView.samza.system", "test");
  configs.put("streams.PageView.source", Base64Serializer.serialize(pageViews));
  configs.put("streams.PageView.partitionCount", String.valueOf(partitionCount));

  configs.put("streams.Profile.samza.system", "test");
  configs.put("streams.Profile.samza.bootstrap", "true");
  configs.put("streams.Profile.source", Base64Serializer.serialize(profiles));
  configs.put("streams.Profile.partitionCount", String.valueOf(partitionCount));

  Config config = new MapConfig(configs);
  final LocalApplicationRunner runner = new LocalApplicationRunner(new StreamTableJoinApp(), config);
  executeRun(runner, config);
  runner.waitForFinish();

  assertEquals(count * partitionCount, StreamTableJoinApp.received.size());
  assertEquals(count * partitionCount, StreamTableJoinApp.joined.size());
  assertTrue(StreamTableJoinApp.joined.get(0) instanceof EnrichedPageView);
}
 

@Test
public void testWithConfigRewriter() throws Exception {
  Map<String, String> configs = getBaseJobConfig(bootstrapUrl(), zkConnect());
  configs.put("streams.PageView.samza.system", "test");
  configs.put("streams.PageView.source", Base64Serializer.serialize(TestTableData.generatePageViews(10)));
  configs.put("streams.PageView.partitionCount", String.valueOf(4));
  configs.put("task.inputs", "test.PageView");
  configs.put("job.config.rewriter.my-rewriter.class", MyConfigRewriter.class.getName());
  configs.put("job.config.rewriters", "my-rewriter");

  Config config = new MapConfig(configs);
  final LocalApplicationRunner runner = new LocalApplicationRunner(new MyTaskApplication(), config);
  executeRun(runner, config);
  runner.waitForFinish();
}
 

/**
 * Executes the application using the local application runner.
 * It takes two required command line arguments
 *  config-factory: a fully {@link org.apache.samza.config.factories.PropertiesConfigFactory} class name
 *  config-path: path to application properties
 *
 * @param args command line arguments
 */
public static void main(String[] args) {
  CommandLine cmdLine = new CommandLine();
  OptionSet options = cmdLine.parser().parse(args);
  Config config = cmdLine.loadConfig(options);

  WikipediaApplication app = new WikipediaApplication();
  LocalApplicationRunner runner = new LocalApplicationRunner(app, config);
  runner.run();
  runner.waitForFinish();
}
 

public static void main(String[] args) {
  CommandLine cmdLine = new CommandLine();
  OptionSet options = cmdLine.parser().parse(args);
  Config config = cmdLine.loadConfig(options);

  AzureApplication app = new AzureApplication();
  LocalApplicationRunner runner = new LocalApplicationRunner(app, config);
  runner.run();

  runner.waitForFinish();
}
 

private void doTestStreamTableJoinRemoteTable(String testName, boolean batchRead, boolean batchWrite) throws Exception {
  final InMemoryWriteFunction writer = new InMemoryWriteFunction(testName);

  batchReads.put(testName, new AtomicInteger());
  batchWrites.put(testName, new AtomicInteger());
  writtenRecords.put(testName, new CopyOnWriteArrayList<>());

  final int count = 16;
  final int batchSize = 4;
  PageView[] pageViews = generatePageViewsWithDistinctKeys(count);
  String profiles = Base64Serializer.serialize(generateProfiles(count));

  int partitionCount = 1;
  Map<String, String> configs = TestLocalTableEndToEnd.getBaseJobConfig(bootstrapUrl(), zkConnect());

  configs.put("streams.PageView.samza.system", "test");
  configs.put("streams.PageView.source", Base64Serializer.serialize(pageViews));
  configs.put("streams.PageView.partitionCount", String.valueOf(partitionCount));
  configs.put("task.max.concurrency", String.valueOf(count));
  configs.put("task.async.commit", String.valueOf(true));

  final RateLimiter readRateLimiter = mock(RateLimiter.class, withSettings().serializable());
  final RateLimiter writeRateLimiter = mock(RateLimiter.class, withSettings().serializable());
  final TableRateLimiter.CreditFunction creditFunction = (k, v, args)->1;
  final StreamApplication app = appDesc -> {
    RemoteTableDescriptor<Integer, Profile> inputTableDesc = new RemoteTableDescriptor<>("profile-table-1");
    inputTableDesc
        .withReadFunction(InMemoryReadFunction.getInMemoryReadFunction(testName, profiles))
        .withRateLimiter(readRateLimiter, creditFunction, null);
    if (batchRead) {
      inputTableDesc.withBatchProvider(new CompactBatchProvider().withMaxBatchSize(batchSize).withMaxBatchDelay(Duration.ofHours(1)));
    }

    // dummy reader
    TableReadFunction readFn = new MyReadFunction();

    RemoteTableDescriptor<Integer, EnrichedPageView> outputTableDesc = new RemoteTableDescriptor<>("enriched-page-view-table-1");
    outputTableDesc
        .withReadFunction(readFn)
        .withWriteFunction(writer)
        .withRateLimiter(writeRateLimiter, creditFunction, creditFunction);
    if (batchWrite) {
      outputTableDesc.withBatchProvider(new CompactBatchProvider().withMaxBatchSize(batchSize).withMaxBatchDelay(Duration.ofHours(1)));
    }

    Table<KV<Integer, EnrichedPageView>> outputTable = appDesc.getTable(outputTableDesc);

    Table<KV<Integer, Profile>> inputTable = appDesc.getTable(inputTableDesc);

    DelegatingSystemDescriptor ksd = new DelegatingSystemDescriptor("test");
    GenericInputDescriptor<PageView> isd = ksd.getInputDescriptor("PageView", new NoOpSerde<>());
    appDesc.getInputStream(isd)
        .map(pv -> new KV<>(pv.getMemberId(), pv))
        .join(inputTable, new PageViewToProfileJoinFunction())
        .map(m -> new KV(m.getMemberId(), m))
        .sendTo(outputTable);
  };

  Config config = new MapConfig(configs);
  final LocalApplicationRunner runner = new LocalApplicationRunner(app, config);
  executeRun(runner, config);

  runner.waitForFinish();
  int numExpected = count * partitionCount;
  Assert.assertEquals(numExpected, writtenRecords.get(testName).size());
  Assert.assertTrue(writtenRecords.get(testName).get(0) instanceof EnrichedPageView);

  if (batchRead) {
    Assert.assertEquals(numExpected / batchSize, batchReads.get(testName).get());
  }
  if (batchWrite) {
    Assert.assertEquals(numExpected / batchSize, batchWrites.get(testName).get());
  }
}
 

private void doTestStreamTableJoinRemoteTable(boolean withCache, boolean defaultCache, boolean withArgs, String testName)
    throws Exception {

  writtenRecords.put(testName, new ArrayList<>());

  int count = 10;
  final PageView[] pageViews = generatePageViews(count);
  final String profiles = Base64Serializer.serialize(generateProfiles(count));

  final int partitionCount = 4;
  final Map<String, String> configs = TestLocalTableEndToEnd.getBaseJobConfig(bootstrapUrl(), zkConnect());
  configs.put("streams.PageView.samza.system", "test");
  configs.put("streams.PageView.source", Base64Serializer.serialize(pageViews));
  configs.put("streams.PageView.partitionCount", String.valueOf(partitionCount));

  final RateLimiter readRateLimiter = mock(RateLimiter.class, withSettings().serializable());
  final TableRateLimiter.CreditFunction creditFunction = (k, v, args) -> 1;
  final StreamApplication app = appDesc -> {

    final RemoteTableDescriptor joinTableDesc =
            new RemoteTableDescriptor<Integer, TestTableData.Profile>("profile-table-1")
        .withReadFunction(InMemoryProfileReadFunction.getInMemoryReadFunction(testName, profiles))
        .withRateLimiter(readRateLimiter, creditFunction, null);

    final RemoteTableDescriptor outputTableDesc =
            new RemoteTableDescriptor<Integer, EnrichedPageView>("enriched-page-view-table-1")
        .withReadFunction(new NoOpTableReadFunction<>())
        .withReadRateLimiterDisabled()
        .withWriteFunction(new InMemoryEnrichedPageViewWriteFunction(testName))
        .withWriteRateLimit(1000);

    final Table<KV<Integer, EnrichedPageView>> outputTable = withCache
        ? getCachingTable(outputTableDesc, defaultCache, appDesc)
        : appDesc.getTable(outputTableDesc);

    final Table<KV<Integer, Profile>> joinTable = withCache
        ? getCachingTable(joinTableDesc, defaultCache, appDesc)
        : appDesc.getTable(joinTableDesc);

    final DelegatingSystemDescriptor ksd = new DelegatingSystemDescriptor("test");
    final GenericInputDescriptor<PageView> isd = ksd.getInputDescriptor("PageView", new NoOpSerde<>());

    if (!withArgs) {
      appDesc.getInputStream(isd)
          .map(pv -> new KV<>(pv.getMemberId(), pv))
          .join(joinTable, new PageViewToProfileJoinFunction())
          .map(m -> new KV(m.getMemberId(), m))
          .sendTo(outputTable);

    } else {
      counters.put(testName, new AtomicInteger());

      final RemoteTableDescriptor counterTableDesc =
          new RemoteTableDescriptor("counter-table-1")
              .withReadFunction(new InMemoryCounterReadFunction(testName))
              .withWriteFunction(new InMemoryCounterWriteFunction(testName))
              .withRateLimiterDisabled();

      final Table counterTable = withCache
          ? getCachingTable(counterTableDesc, defaultCache, appDesc)
          : appDesc.getTable(counterTableDesc);

      final String counterTableName = ((TableImpl) counterTable).getTableId();
      appDesc.getInputStream(isd)
          .map(new TestReadWriteMapFunction(counterTableName))
          .map(pv -> new KV<>(pv.getMemberId(), pv))
          .join(joinTable, new PageViewToProfileJoinFunction(), true)
          .map(m -> new KV(m.getMemberId(), m))
          .sendTo(outputTable, true);
    }
  };

  final Config config = new MapConfig(configs);
  final LocalApplicationRunner runner = new LocalApplicationRunner(app, config);
  executeRun(runner, config);
  runner.waitForFinish();

  final int numExpected = count * partitionCount;
  Assert.assertEquals(numExpected, writtenRecords.get(testName).size());
  Assert.assertTrue(writtenRecords.get(testName).get(0) instanceof EnrichedPageView);
  if (!withArgs) {
    writtenRecords.get(testName).forEach(epv -> Assert.assertFalse(epv.company.contains("-")));
  } else {
    writtenRecords.get(testName).forEach(epv -> Assert.assertTrue(epv.company.endsWith("-r-w")));
    Assert.assertEquals(numExpected, counters.get(testName).get());
  }
}
 

@Test
public void testDualStreamTableJoin() throws Exception {

  int count = 10;
  PageView[] pageViews = TestTableData.generatePageViews(count);
  Profile[] profiles = TestTableData.generateProfiles(count);

  int partitionCount = 4;
  Map<String, String> configs = getBaseJobConfig(bootstrapUrl(), zkConnect());

  configs.put("streams.Profile1.samza.system", "test");
  configs.put("streams.Profile1.source", Base64Serializer.serialize(profiles));
  configs.put("streams.Profile1.samza.bootstrap", "true");
  configs.put("streams.Profile1.partitionCount", String.valueOf(partitionCount));

  configs.put("streams.Profile2.samza.system", "test");
  configs.put("streams.Profile2.source", Base64Serializer.serialize(profiles));
  configs.put("streams.Profile2.samza.bootstrap", "true");
  configs.put("streams.Profile2.partitionCount", String.valueOf(partitionCount));

  configs.put("streams.PageView1.samza.system", "test");
  configs.put("streams.PageView1.source", Base64Serializer.serialize(pageViews));
  configs.put("streams.PageView1.partitionCount", String.valueOf(partitionCount));

  configs.put("streams.PageView2.samza.system", "test");
  configs.put("streams.PageView2.source", Base64Serializer.serialize(pageViews));
  configs.put("streams.PageView2.partitionCount", String.valueOf(partitionCount));

  Config config = new MapConfig(configs);
  final LocalApplicationRunner runner = new LocalApplicationRunner(new DualStreamTableJoinApp(), config);
  executeRun(runner, config);
  runner.waitForFinish();

  assertEquals(count * partitionCount, DualStreamTableJoinApp.sentToProfileTable1.size());
  assertEquals(count * partitionCount, DualStreamTableJoinApp.sentToProfileTable2.size());

  assertEquals(count * partitionCount, DualStreamTableJoinApp.joinedPageViews1.size());
  assertEquals(count * partitionCount, DualStreamTableJoinApp.joinedPageViews2.size());
  assertTrue(DualStreamTableJoinApp.joinedPageViews1.get(0) instanceof EnrichedPageView);
  assertTrue(DualStreamTableJoinApp.joinedPageViews2.get(0) instanceof EnrichedPageView);
}
 

@Test
public void testEndToEnd() throws Exception {

  Bucket inputBucket = cluster.openBucket(inputBucketName);
  inputBucket.upsert(ByteArrayDocument.create("Alice", "20".getBytes()));
  inputBucket.upsert(ByteArrayDocument.create("Bob", "30".getBytes()));
  inputBucket.upsert(ByteArrayDocument.create("Chris", "40".getBytes()));
  inputBucket.upsert(ByteArrayDocument.create("David", "50".getBytes()));
  inputBucket.close();

  String[] users = new String[]{"Alice", "Bob", "Chris", "David"};

  int partitionCount = 1;
  Map<String, String> configs = TestLocalTableEndToEnd.getBaseJobConfig(bootstrapUrl(), zkConnect());

  configs.put("streams.User.samza.system", "test");
  configs.put("streams.User.source", Base64Serializer.serialize(users));
  configs.put("streams.User.partitionCount", String.valueOf(partitionCount));
  Config config = new MapConfig(configs);

  final StreamApplication app = appDesc -> {
    DelegatingSystemDescriptor inputSystemDescriptor = new DelegatingSystemDescriptor("test");
    GenericInputDescriptor<String> inputDescriptor =
        inputSystemDescriptor.getInputDescriptor("User", new NoOpSerde<>());

    CouchbaseTableReadFunction<String> readFunction = new CouchbaseTableReadFunction<>(inputBucketName,
            String.class, "couchbase://127.0.0.1")
        .withBootstrapCarrierDirectPort(couchbaseMock.getCarrierPort(inputBucketName))
        .withBootstrapHttpDirectPort(couchbaseMock.getHttpPort())
        .withSerde(new StringSerde());

    CouchbaseTableWriteFunction<JsonObject> writeFunction = new CouchbaseTableWriteFunction<>(outputBucketName,
            JsonObject.class, "couchbase://127.0.0.1")
        .withBootstrapCarrierDirectPort(couchbaseMock.getCarrierPort(outputBucketName))
        .withBootstrapHttpDirectPort(couchbaseMock.getHttpPort());

    RemoteTableDescriptor inputTableDesc = new RemoteTableDescriptor<String, String>("input-table")
        .withReadFunction(readFunction)
        .withRateLimiterDisabled();
    Table<KV<String, String>> inputTable = appDesc.getTable(inputTableDesc);

    RemoteTableDescriptor outputTableDesc = new RemoteTableDescriptor<String, JsonObject>("output-table")
        .withReadFunction(new NoOpTableReadFunction<>())
        .withWriteFunction(writeFunction)
        .withRateLimiterDisabled();
    Table<KV<String, JsonObject>> outputTable = appDesc.getTable(outputTableDesc);

    appDesc.getInputStream(inputDescriptor)
        .map(k -> KV.of(k, k))
        .join(inputTable, new JoinFunction())
        .sendTo(outputTable);
  };

  final LocalApplicationRunner runner = new LocalApplicationRunner(app, config);
  executeRun(runner, config);
  runner.waitForFinish();

  Bucket outputBucket = cluster.openBucket(outputBucketName);
  Assert.assertEquals("{\"name\":\"Alice\",\"age\":\"20\"}", outputBucket.get("Alice").content().toString());
  Assert.assertEquals("{\"name\":\"Bob\",\"age\":\"30\"}", outputBucket.get("Bob").content().toString());
  Assert.assertEquals("{\"name\":\"Chris\",\"age\":\"40\"}", outputBucket.get("Chris").content().toString());
  Assert.assertEquals("{\"name\":\"David\",\"age\":\"50\"}", outputBucket.get("David").content().toString());
  outputBucket.close();
}