reporters = new HashMap<>();
- for (int i = 0; i < options.getMetricsReporters().size(); i++) {
- final String name = "beam-metrics-reporter-" + i;
- final MetricsReporter reporter = options.getMetricsReporters().get(i);
-
- reporters.put(name, (MetricsReporterFactory) (nm, processorId, config) -> reporter);
- LOG.info("{}: {}", name, reporter.getClass().getName());
- }
- return reporters;
- } else {
- return Collections.emptyMap();
- }
- }
-
- private ApplicationRunner runSamzaApp(StreamApplication app, Config config) {
-
- final ApplicationRunner runner = ApplicationRunners.getApplicationRunner(app, config);
-
- ExternalContext externalContext = null;
- if (listener != null) {
- externalContext = listener.onStart();
- }
-
- runner.run(externalContext);
-
- if (listener != null
- && options.getSamzaExecutionEnvironment() == SamzaExecutionEnvironment.YARN) {
- listener.onSubmit();
- }
-
- return runner;
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaRunnerOverrideConfigs.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaRunnerOverrideConfigs.java
deleted file mode 100644
index 4c5fa432ca82..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaRunnerOverrideConfigs.java
+++ /dev/null
@@ -1,85 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza;
-
-import java.time.Duration;
-
-// TODO: can we get rid of this class? Right now the SamzaPipelineOptionsValidator would force
-// the pipeline option to be the type SamzaPipelineOption. Ideally, we should be able to keep
-// passing SamzaPortablePipelineOption. Alternative, we could merge portable and non-portable
-// pipeline option.
-/** A helper class for holding all the beam runner specific samza configs. */
-@SuppressWarnings({
- "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
-public class SamzaRunnerOverrideConfigs {
- public static final String BEAM_RUNNER_CONFIG_PREFIX = "beam.override.";
- // whether the job is in portable mode
- public static final String IS_PORTABLE_MODE = BEAM_RUNNER_CONFIG_PREFIX + "portable";
- // for portable mode only: port number for fn control api
- public static final String FN_CONTROL_PORT = BEAM_RUNNER_CONFIG_PREFIX + "control.port";
- // timeout for waiting for control client to connect
- public static final String CONTROL_CLIENT_MAX_WAIT_TIME_MS = "controL.wait.time.ms";
- public static final long DEFAULT_CONTROL_CLIENT_MAX_WAIT_TIME_MS =
- Duration.ofMinutes(2).toMillis();
- public static final String FS_TOKEN_PATH = BEAM_RUNNER_CONFIG_PREFIX + "fs.token.path";
- public static final String DEFAULT_FS_TOKEN_PATH = null;
-
- private static boolean containsKey(SamzaPipelineOptions options, String configKey) {
- if (options == null || options.getConfigOverride() == null) {
- return false;
- }
- return options.getConfigOverride().containsKey(configKey);
- }
-
- /** Whether the job is in portable mode based on the config override in the pipeline options. */
- public static boolean isPortableMode(SamzaPipelineOptions options) {
- if (containsKey(options, IS_PORTABLE_MODE)) {
- return options.getConfigOverride().get(IS_PORTABLE_MODE).equals(String.valueOf(true));
- } else {
- return false;
- }
- }
-
- /** Get fn control port number based on the config override in the pipeline options. */
- public static int getFnControlPort(SamzaPipelineOptions options) {
- if (containsKey(options, FN_CONTROL_PORT)) {
- return Integer.parseInt(options.getConfigOverride().get(FN_CONTROL_PORT));
- } else {
- return -1;
- }
- }
-
- /** Get max wait time for control client connection. */
- public static long getControlClientWaitTimeoutMs(SamzaPipelineOptions options) {
- if (containsKey(options, CONTROL_CLIENT_MAX_WAIT_TIME_MS)) {
- return Long.parseLong(options.getConfigOverride().get(CONTROL_CLIENT_MAX_WAIT_TIME_MS));
- } else {
- return DEFAULT_CONTROL_CLIENT_MAX_WAIT_TIME_MS;
- }
- }
-
- /** Get fs token path for portable mode. */
- public static String getFsTokenPath(SamzaPipelineOptions options) {
- if (containsKey(options, FS_TOKEN_PATH)) {
- return options.getConfigOverride().get(FS_TOKEN_PATH);
- } else {
- return DEFAULT_FS_TOKEN_PATH;
- }
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaRunnerRegistrar.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaRunnerRegistrar.java
deleted file mode 100644
index 102838975d13..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaRunnerRegistrar.java
+++ /dev/null
@@ -1,53 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza;
-
-import com.google.auto.service.AutoService;
-import org.apache.beam.sdk.PipelineRunner;
-import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.options.PipelineOptionsRegistrar;
-import org.apache.beam.sdk.runners.PipelineRunnerRegistrar;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
-
-/**
- * AutoService registrar - will register SamzaRunner and SamzaOptions as possible pipeline runner
- * services.
- *
- * It ends up in META-INF/services and gets picked up by Beam.
- */
-public class SamzaRunnerRegistrar {
- private SamzaRunnerRegistrar() {}
-
- /** Pipeline runner registrar. */
- @AutoService(PipelineRunnerRegistrar.class)
- public static class Runner implements PipelineRunnerRegistrar {
- @Override
- public Iterable>> getPipelineRunners() {
- return ImmutableList.of(SamzaRunner.class, TestSamzaRunner.class);
- }
- }
-
- /** Pipeline options registrar. */
- @AutoService(PipelineOptionsRegistrar.class)
- public static class Options implements PipelineOptionsRegistrar {
- @Override
- public Iterable> getPipelineOptions() {
- return ImmutableList.of(SamzaPipelineOptions.class);
- }
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/TestSamzaRunner.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/TestSamzaRunner.java
deleted file mode 100644
index 810fc0c983f7..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/TestSamzaRunner.java
+++ /dev/null
@@ -1,109 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza;
-
-import static org.apache.samza.config.JobConfig.JOB_JMX_ENABLED;
-import static org.apache.samza.config.JobConfig.JOB_LOGGED_STORE_BASE_DIR;
-import static org.apache.samza.config.JobConfig.JOB_NON_LOGGED_STORE_BASE_DIR;
-
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
-import java.io.File;
-import java.nio.file.Files;
-import java.util.HashMap;
-import java.util.Map;
-import org.apache.beam.runners.samza.translation.ConfigBuilder;
-import org.apache.beam.sdk.Pipeline;
-import org.apache.beam.sdk.PipelineResult;
-import org.apache.beam.sdk.PipelineRunner;
-import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.options.PipelineOptionsValidator;
-import org.apache.commons.io.FileUtils;
-
-/** Test {@link SamzaRunner}. */
-public class TestSamzaRunner extends PipelineRunner {
-
- private final SamzaRunner delegate;
- private final File storeDir;
-
- public static TestSamzaRunner fromOptions(PipelineOptions options) {
- return new TestSamzaRunner(options);
- }
-
- public static SamzaPipelineOptions createSamzaPipelineOptions(
- PipelineOptions options, File storeDir) {
- try {
- final SamzaPipelineOptions samzaOptions =
- PipelineOptionsValidator.validate(SamzaPipelineOptions.class, options);
- final Map config = new HashMap<>(ConfigBuilder.localRunConfig());
- config.put(JOB_LOGGED_STORE_BASE_DIR, storeDir.getAbsolutePath());
- config.put(JOB_NON_LOGGED_STORE_BASE_DIR, storeDir.getAbsolutePath());
- config.put(JOB_JMX_ENABLED, "false");
-
- if (samzaOptions.getConfigOverride() != null) {
- config.putAll(samzaOptions.getConfigOverride());
- }
- samzaOptions.setConfigOverride(config);
- return samzaOptions;
- } catch (Exception e) {
- throw new RuntimeException(e);
- }
- }
-
- private static File createStoreDir() {
- try {
- return Files.createTempDirectory("beam-samza-test").toFile();
- } catch (Exception e) {
- throw new RuntimeException(e);
- }
- }
-
- public TestSamzaRunner(PipelineOptions options) {
- this.storeDir = createStoreDir();
- this.delegate = SamzaRunner.fromOptions(createSamzaPipelineOptions(options, storeDir));
- }
-
- @Override
- @SuppressFBWarnings(value = "DE_MIGHT_IGNORE")
- public PipelineResult run(Pipeline pipeline) {
- try {
- final PipelineResult result = delegate.run(pipeline);
- result.waitUntilFinish();
-
- return result;
- } catch (Throwable t) {
- // Search for AssertionError. If present use it as the cause of the pipeline failure.
- Throwable current = t;
-
- while (current != null) {
- if (current instanceof AssertionError) {
- throw (AssertionError) current;
- }
- current = current.getCause();
- }
-
- throw t;
- } finally {
- try {
- // delete the store folder
- FileUtils.deleteDirectory(storeDir);
- } catch (Exception ignore) {
- // Ignore
- }
- }
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/adapter/BoundedSourceSystem.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/adapter/BoundedSourceSystem.java
deleted file mode 100644
index 92c9eea4293a..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/adapter/BoundedSourceSystem.java
+++ /dev/null
@@ -1,449 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.adapter;
-
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.LinkedBlockingQueue;
-import java.util.concurrent.Semaphore;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.function.Function;
-import java.util.stream.Collectors;
-import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
-import org.apache.beam.runners.core.serialization.Base64Serializer;
-import org.apache.beam.runners.samza.SamzaPipelineOptions;
-import org.apache.beam.runners.samza.metrics.FnWithMetricsWrapper;
-import org.apache.beam.runners.samza.metrics.SamzaMetricsContainer;
-import org.apache.beam.runners.samza.runtime.OpMessage;
-import org.apache.beam.sdk.io.BoundedSource;
-import org.apache.beam.sdk.io.BoundedSource.BoundedReader;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.values.WindowedValue;
-import org.apache.beam.sdk.values.WindowedValues;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
-import org.apache.samza.Partition;
-import org.apache.samza.SamzaException;
-import org.apache.samza.config.Config;
-import org.apache.samza.metrics.MetricsRegistry;
-import org.apache.samza.metrics.MetricsRegistryMap;
-import org.apache.samza.system.IncomingMessageEnvelope;
-import org.apache.samza.system.SystemAdmin;
-import org.apache.samza.system.SystemConsumer;
-import org.apache.samza.system.SystemFactory;
-import org.apache.samza.system.SystemProducer;
-import org.apache.samza.system.SystemStreamMetadata;
-import org.apache.samza.system.SystemStreamMetadata.SystemStreamPartitionMetadata;
-import org.apache.samza.system.SystemStreamPartition;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * A Samza system that supports reading from a Beam {@link BoundedSource}. The source is treated as
- * though it has a single partition and does not support checkpointing via a changelog stream. If
- * the job is restarted the bounded source will be consumed from the beginning.
- */
-// TODO: instrumentation for the consumer
-@SuppressWarnings({
- "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
-public class BoundedSourceSystem {
- private static final Logger LOG = LoggerFactory.getLogger(BoundedSourceSystem.class);
-
- private static List> split(
- BoundedSource source, SamzaPipelineOptions pipelineOptions) throws Exception {
- final int numSplits = pipelineOptions.getMaxSourceParallelism();
- if (numSplits > 1) {
- final long estimatedSize = source.getEstimatedSizeBytes(pipelineOptions);
- // calculate the size of each split, rounded up to the ceiling.
- final long bundleSize = (estimatedSize + numSplits - 1) / numSplits;
- @SuppressWarnings("unchecked")
- final List> splits =
- (List>) source.split(bundleSize, pipelineOptions);
- // Need the empty check here because Samza doesn't handle empty partition well
- if (!splits.isEmpty()) {
- return splits;
- }
- }
- return Collections.singletonList(source);
- }
-
- /** A {@link SystemAdmin} for {@link BoundedSourceSystem}. */
- public static class Admin implements SystemAdmin {
- private final BoundedSource source;
- private final SamzaPipelineOptions pipelineOptions;
-
- public Admin(BoundedSource source, SamzaPipelineOptions pipelineOptions) {
- this.source = source;
- this.pipelineOptions = pipelineOptions;
- }
-
- @Override
- public Map getOffsetsAfter(
- Map offsets) {
- // BEAM checkpoints the next offset so here we just need to return the map itself
- return offsets;
- }
-
- @Override
- public Map getSystemStreamMetadata(Set streamNames) {
- return streamNames.stream()
- .collect(
- Collectors.toMap(
- Function.identity(),
- streamName -> {
- try {
- List> splits = split(source, pipelineOptions);
- final Map partitionMetaData =
- new HashMap<>();
- // we assume that the generated splits are stable,
- // this is necessary so that the mapping of partition to source is correct
- // in each container.
- for (int i = 0; i < splits.size(); i++) {
- partitionMetaData.put(
- new Partition(i), new SystemStreamPartitionMetadata(null, null, null));
- }
- return new SystemStreamMetadata(streamName, partitionMetaData);
- } catch (Exception e) {
- throw new SamzaException("Fail to read stream metadata", e);
- }
- }));
- }
-
- @Override
- public Integer offsetComparator(String offset1, String offset2) {
- if (offset1 == null) {
- return offset2 == null ? 0 : -1;
- }
-
- if (offset2 == null) {
- return 1;
- }
-
- return Long.valueOf(offset1).compareTo(Long.valueOf(offset2));
- }
- }
-
- /**
- * A {@link SystemConsumer} for a {@link BoundedSource}. See {@link BoundedSourceSystem} for more
- * details.
- */
- public static class Consumer implements SystemConsumer {
- private static final Logger LOG = LoggerFactory.getLogger(Consumer.class);
- private static final AtomicInteger NEXT_ID = new AtomicInteger();
-
- private final List> splits;
- private final SamzaPipelineOptions pipelineOptions;
- private final Map, SystemStreamPartition> readerToSsp = new HashMap<>();
- private final SamzaMetricsContainer metricsContainer;
- private final String stepName;
-
- private ReaderTask readerTask;
-
- Consumer(
- BoundedSource source,
- SamzaPipelineOptions pipelineOptions,
- SamzaMetricsContainer metricsContainer,
- String stepName) {
- try {
- splits = split(source, pipelineOptions);
- } catch (Exception e) {
- throw new SamzaException("Fail to split source", e);
- }
- this.pipelineOptions = pipelineOptions;
- this.metricsContainer = metricsContainer;
- this.stepName = stepName;
- }
-
- @Override
- public void start() {
- if (this.readerToSsp.isEmpty()) {
- throw new IllegalArgumentException(
- "Attempted to call start without assigned system stream partitions");
- }
-
- final int capacity = pipelineOptions.getSystemBufferSize();
- final FnWithMetricsWrapper metricsWrapper =
- pipelineOptions.getEnableMetrics()
- ? new FnWithMetricsWrapper(metricsContainer, stepName)
- : null;
- readerTask = new ReaderTask<>(readerToSsp, capacity, metricsWrapper);
- final Thread thread =
- new Thread(readerTask, "bounded-source-system-consumer-" + NEXT_ID.getAndIncrement());
- thread.start();
- }
-
- @Override
- public void stop() {
- // NOTE: this is not a blocking shutdown
- if (readerTask != null) {
- readerTask.stop();
- }
- }
-
- @Override
- public void register(SystemStreamPartition ssp, String offset) {
- final int partitionId = ssp.getPartition().getPartitionId();
- try {
- final BoundedReader reader = splits.get(partitionId).createReader(pipelineOptions);
- readerToSsp.put(reader, ssp);
- } catch (Exception e) {
- throw new SamzaException("Error while creating source reader for ssp: " + ssp, e);
- }
- }
-
- @Override
- public Map> poll(
- Set systemStreamPartitions, long timeout)
- throws InterruptedException {
- assert !readerToSsp.isEmpty(); // start should be called before poll
-
- final Map> envelopes = new HashMap<>();
- for (SystemStreamPartition ssp : systemStreamPartitions) {
- envelopes.put(ssp, readerTask.getNextMessages(ssp, timeout));
- }
- return envelopes;
- }
-
- private static class ReaderTask implements Runnable {
- private final Map, SystemStreamPartition> readerToSsp;
- private final Map> queues;
- private final Semaphore available;
- private final FnWithMetricsWrapper metricsWrapper;
-
- // NOTE: we do not support recovery with a bounded source (we restart from the beginning),
- // so we do not need to have a way to tie an offset to a position in the bounded source.
- private long offset;
- private volatile Thread readerThread;
- private volatile boolean stopInvoked = false;
- private volatile Exception lastException;
-
- private ReaderTask(
- Map, SystemStreamPartition> readerToSsp,
- int capacity,
- FnWithMetricsWrapper metricsWrapper) {
- this.readerToSsp = readerToSsp;
- this.available = new Semaphore(capacity);
- this.metricsWrapper = metricsWrapper;
-
- final Map> qs =
- new HashMap<>();
- readerToSsp.values().forEach(ssp -> qs.put(ssp, new LinkedBlockingQueue<>()));
- this.queues = ImmutableMap.copyOf(qs);
- }
-
- @Override
- public void run() {
- readerThread = Thread.currentThread();
-
- final Set> availableReaders = new HashSet<>(readerToSsp.keySet());
- try {
- for (BoundedReader reader : readerToSsp.keySet()) {
- boolean hasData = invoke(reader::start);
- if (hasData) {
- enqueueMessage(reader);
- } else {
- enqueueMaxWatermarkAndEndOfStream(reader);
- reader.close();
- availableReaders.remove(reader);
- }
- }
-
- while (!stopInvoked && !availableReaders.isEmpty()) {
- final Iterator> iter = availableReaders.iterator();
- while (iter.hasNext()) {
- final BoundedReader reader = iter.next();
- final boolean hasData = invoke(reader::advance);
- if (hasData) {
- enqueueMessage(reader);
- } else {
- enqueueMaxWatermarkAndEndOfStream(reader);
- reader.close();
- iter.remove();
- }
- }
- }
- } catch (InterruptedException e) {
- // We use an interrupt to wake the reader from a blocking read under normal termination,
- // so ignore it here.
- } catch (Exception e) {
- setError(e);
- } finally {
- availableReaders.forEach(
- reader -> {
- try {
- reader.close();
- } catch (IOException e) {
- LOG.error(
- "Reader task failed to close reader for ssp {}", readerToSsp.get(reader), e);
- }
- });
- }
- }
-
- private X invoke(FnWithMetricsWrapper.SupplierWithException fn) throws Exception {
- if (metricsWrapper != null) {
- return metricsWrapper.wrap(fn, true);
- } else {
- return fn.get();
- }
- }
-
- private void enqueueMessage(BoundedReader reader) throws InterruptedException {
- final T value = reader.getCurrent();
- final WindowedValue windowedValue =
- WindowedValues.timestampedValueInGlobalWindow(value, reader.getCurrentTimestamp());
- final SystemStreamPartition ssp = readerToSsp.get(reader);
- final IncomingMessageEnvelope envelope =
- new IncomingMessageEnvelope(
- ssp, Long.toString(offset++), null, OpMessage.ofElement(windowedValue));
-
- available.acquire();
- queues.get(ssp).put(envelope);
- }
-
- private void enqueueMaxWatermarkAndEndOfStream(BoundedReader reader) {
- final SystemStreamPartition ssp = readerToSsp.get(reader);
- // Send the max watermark to force completion of any open windows.
- final IncomingMessageEnvelope watermarkEnvelope =
- IncomingMessageEnvelope.buildWatermarkEnvelope(
- ssp, BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis());
- enqueueUninterruptibly(watermarkEnvelope);
-
- final IncomingMessageEnvelope endOfStreamEnvelope =
- IncomingMessageEnvelope.buildEndOfStreamEnvelope(ssp);
- enqueueUninterruptibly(endOfStreamEnvelope);
- }
-
- private void stop() {
- stopInvoked = true;
-
- final Thread readerThread = this.readerThread;
- if (readerThread != null) {
- readerThread.interrupt();
- }
- }
-
- private List getNextMessages(
- SystemStreamPartition ssp, long timeoutMillis) throws InterruptedException {
- if (lastException != null) {
- throw new RuntimeException(lastException);
- }
-
- final List envelopes = new ArrayList<>();
- final BlockingQueue queue = queues.get(ssp);
- final IncomingMessageEnvelope envelope = queue.poll(timeoutMillis, TimeUnit.MILLISECONDS);
-
- if (envelope != null) {
- envelopes.add(envelope);
- queue.drainTo(envelopes);
- }
-
- available.release(envelopes.size());
-
- if (lastException != null) {
- throw new RuntimeException(lastException);
- }
-
- return envelopes;
- }
-
- private void setError(Exception exception) {
- this.lastException = exception;
- // A dummy message used to force the consumer to wake up immediately and check the
- // lastException field, which will be populated.
- readerToSsp
- .values()
- .forEach(
- ssp -> {
- final IncomingMessageEnvelope checkLastExceptionEvelope =
- new IncomingMessageEnvelope(ssp, null, null, null);
- enqueueUninterruptibly(checkLastExceptionEvelope);
- });
- }
-
- private void enqueueUninterruptibly(IncomingMessageEnvelope envelope) {
- final BlockingQueue queue =
- queues.get(envelope.getSystemStreamPartition());
- while (true) {
- try {
- queue.put(envelope);
- return;
- } catch (InterruptedException e) {
- // Some events require that we post an envelope to the queue even if the interrupt
- // flag was set (i.e. during a call to stop) to ensure that the consumer properly
- // shuts down. Consequently, if we receive an interrupt here we ignore it and retry
- // the put operation.
- }
- }
- }
- }
- }
-
- /**
- * A {@link SystemFactory} that produces a {@link BoundedSourceSystem} for a particular {@link
- * BoundedSource} registered in {@link Config}.
- */
- public static class Factory implements SystemFactory {
- @Override
- public SystemConsumer getConsumer(String systemName, Config config, MetricsRegistry registry) {
- final String streamPrefix = "systems." + systemName;
- final Config scopedConfig = config.subset(streamPrefix + ".", true);
-
- return new Consumer(
- getBoundedSource(scopedConfig),
- getPipelineOptions(config),
- new SamzaMetricsContainer((MetricsRegistryMap) registry),
- scopedConfig.get("stepName"));
- }
-
- @Override
- public SystemProducer getProducer(String systemName, Config config, MetricsRegistry registry) {
- LOG.info("System {} does not have producer.", systemName);
- return null;
- }
-
- @Override
- public SystemAdmin getAdmin(String systemName, Config config) {
- final Config scopedConfig = config.subset("systems." + systemName + ".", true);
- return new Admin(getBoundedSource(scopedConfig), getPipelineOptions(config));
- }
-
- private static BoundedSource getBoundedSource(Config config) {
- @SuppressWarnings("unchecked")
- final BoundedSource source =
- Base64Serializer.deserializeUnchecked(config.get("source"), BoundedSource.class);
- return source;
- }
-
- private static SamzaPipelineOptions getPipelineOptions(Config config) {
- return Base64Serializer.deserializeUnchecked(
- config.get("beamPipelineOptions"), SerializablePipelineOptions.class)
- .get()
- .as(SamzaPipelineOptions.class);
- }
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/adapter/UnboundedSourceSystem.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/adapter/UnboundedSourceSystem.java
deleted file mode 100644
index ffab2ff59ce5..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/adapter/UnboundedSourceSystem.java
+++ /dev/null
@@ -1,533 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.adapter;
-
-import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Base64;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.LinkedBlockingQueue;
-import java.util.concurrent.Semaphore;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.function.Function;
-import java.util.stream.Collectors;
-import org.apache.beam.repackaged.core.org.apache.commons.lang3.StringUtils;
-import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
-import org.apache.beam.runners.core.serialization.Base64Serializer;
-import org.apache.beam.runners.samza.SamzaPipelineOptions;
-import org.apache.beam.runners.samza.metrics.FnWithMetricsWrapper;
-import org.apache.beam.runners.samza.metrics.SamzaMetricsContainer;
-import org.apache.beam.runners.samza.runtime.OpMessage;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.io.UnboundedSource;
-import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark;
-import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.values.WindowedValue;
-import org.apache.beam.sdk.values.WindowedValues;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
-import org.apache.samza.Partition;
-import org.apache.samza.SamzaException;
-import org.apache.samza.config.Config;
-import org.apache.samza.metrics.MetricsRegistry;
-import org.apache.samza.metrics.MetricsRegistryMap;
-import org.apache.samza.system.IncomingMessageEnvelope;
-import org.apache.samza.system.SystemAdmin;
-import org.apache.samza.system.SystemConsumer;
-import org.apache.samza.system.SystemFactory;
-import org.apache.samza.system.SystemProducer;
-import org.apache.samza.system.SystemStreamMetadata;
-import org.apache.samza.system.SystemStreamMetadata.SystemStreamPartitionMetadata;
-import org.apache.samza.system.SystemStreamPartition;
-import org.joda.time.Instant;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * A Samza system that supports reading from a Beam {@link UnboundedSource}. The source is split
- * into partitions. Samza creates the job model by assigning partitions to Samza tasks.
- */
-@SuppressWarnings({
- "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
- "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
-public class UnboundedSourceSystem {
- private static final Logger LOG = LoggerFactory.getLogger(UnboundedSourceSystem.class);
-
- // A dummy message used to force the consumer to wake up immediately and check the
- // lastException field, which will be populated.
- private static final IncomingMessageEnvelope CHECK_LAST_EXCEPTION_ENVELOPE =
- new IncomingMessageEnvelope(null, null, null, null);
-
- /**
- * For better parallelism in Samza, we need to configure a large split number for {@link
- * UnboundedSource} like Kafka. This will most likely make each split contain a single partition,
- * and be assigned to a Samza task. A large split number is safe since the actual split is bounded
- * by the number of source partitions.
- */
- private static
- List> split(
- UnboundedSource source, SamzaPipelineOptions pipelineOptions)
- throws Exception {
- final int numSplits = pipelineOptions.getMaxSourceParallelism();
- if (numSplits > 1) {
- @SuppressWarnings("unchecked")
- final List> splits =
- (List>) source.split(numSplits, pipelineOptions);
- // Need the empty check here because Samza doesn't handle empty partition well
- if (!splits.isEmpty()) {
- return splits;
- }
- }
- return Collections.singletonList(source);
- }
-
- /** A {@link SystemAdmin} for {@link UnboundedSourceSystem}. */
- public static class Admin implements SystemAdmin {
- private final UnboundedSource source;
- private final SamzaPipelineOptions pipelineOptions;
-
- public Admin(UnboundedSource source, SamzaPipelineOptions pipelineOptions) {
- this.source = source;
- this.pipelineOptions = pipelineOptions;
- }
-
- @Override
- public Map getOffsetsAfter(
- Map offsets) {
- // BEAM checkpoints the next offset so here we just need to return the map itself
- return offsets;
- }
-
- @Override
- public Map getSystemStreamMetadata(Set streamNames) {
- return streamNames.stream()
- .collect(
- Collectors.toMap(
- Function.identity(),
- streamName -> {
- try {
- final List> splits =
- split(source, pipelineOptions);
- final Map partitionMetaData =
- new HashMap<>();
- // we assume that the generated splits are stable,
- // this is necessary so that the mapping of partition to source is correct
- // in each container.
- for (int i = 0; i < splits.size(); i++) {
- partitionMetaData.put(
- new Partition(i), new SystemStreamPartitionMetadata(null, null, null));
- }
- return new SystemStreamMetadata(streamName, partitionMetaData);
- } catch (Exception e) {
- throw new SamzaException("Fail to read stream metadata", e);
- }
- }));
- }
-
- @Override
- public Integer offsetComparator(String offset1, String offset2) {
- // BEAM will fetch the exact offset. So we don't need to compare them.
- // Return null indicating it's caught up.
- return null;
- }
- }
-
- /**
- * A {@link SystemConsumer} for a {@link UnboundedSource}. See {@link UnboundedSourceSystem} for
- * more details.
- */
- public static class Consumer
- implements SystemConsumer {
- private static final Logger LOG = LoggerFactory.getLogger(Consumer.class);
-
- private static final AtomicInteger NEXT_ID = new AtomicInteger();
-
- private final Coder checkpointMarkCoder;
- private final List> splits;
- private final SamzaPipelineOptions pipelineOptions;
- private final Map readerToSsp = new HashMap<>();
- private final SamzaMetricsContainer metricsContainer;
- private final String stepName;
-
- private ReaderTask readerTask;
-
- Consumer(
- UnboundedSource source,
- SamzaPipelineOptions pipelineOptions,
- SamzaMetricsContainer metricsContainer,
- String stepName) {
- try {
- this.splits = split(source, pipelineOptions);
- } catch (Exception e) {
- throw new SamzaException("Fail to split source", e);
- }
- this.checkpointMarkCoder = source.getCheckpointMarkCoder();
- this.pipelineOptions = pipelineOptions;
- this.metricsContainer = metricsContainer;
- this.stepName = stepName;
- }
-
- @Override
- public void start() {
- if (this.readerToSsp.isEmpty()) {
- throw new IllegalArgumentException(
- "Attempted to call start without assigned system stream partitions");
- }
-
- final FnWithMetricsWrapper metricsWrapper =
- pipelineOptions.getEnableMetrics()
- ? new FnWithMetricsWrapper(metricsContainer, stepName)
- : null;
- readerTask =
- new ReaderTask<>(
- readerToSsp,
- checkpointMarkCoder,
- pipelineOptions.getSystemBufferSize(),
- pipelineOptions.getWatermarkInterval(),
- metricsWrapper);
- final Thread thread =
- new Thread(readerTask, "unbounded-source-system-consumer-" + NEXT_ID.getAndIncrement());
- thread.start();
- }
-
- @Override
- public void stop() {
- // NOTE: this is not a blocking shutdown
- readerTask.stop();
- }
-
- @Override
- public void register(SystemStreamPartition ssp, String offset) {
- CheckpointMarkT checkpoint = null;
- if (StringUtils.isNoneEmpty(offset)) {
- final byte[] offsetBytes = Base64.getDecoder().decode(offset);
- final ByteArrayInputStream bais = new ByteArrayInputStream(offsetBytes);
- try {
- checkpoint = checkpointMarkCoder.decode(bais);
- } catch (Exception e) {
- throw new SamzaException("Error in decode offset", e);
- }
- }
-
- // Create unbounded reader with checkpoint
- final int partitionId = ssp.getPartition().getPartitionId();
- try {
- final UnboundedReader reader =
- splits.get(partitionId).createReader(pipelineOptions, checkpoint);
- readerToSsp.put(reader, ssp);
- } catch (Exception e) {
- throw new SamzaException("Error while creating source reader for ssp: " + ssp, e);
- }
- }
-
- @Override
- public Map> poll(
- Set systemStreamPartitions, long timeout)
- throws InterruptedException {
- assert !readerToSsp.isEmpty(); // start should be called before poll
-
- final Map> envelopes = new HashMap<>();
- for (SystemStreamPartition ssp : systemStreamPartitions) {
- envelopes.put(ssp, readerTask.getNextMessages(ssp, timeout));
- }
- return envelopes;
- }
-
- private static class ReaderTask implements Runnable {
- private final Map readerToSsp;
- private final List readers;
- private final Coder checkpointMarkCoder;
- private final Map currentWatermarks = new HashMap<>();
- private final Map> queues;
- private final long watermarkInterval;
- private final Semaphore available;
- private final FnWithMetricsWrapper metricsWrapper;
-
- private volatile boolean running;
- private volatile Exception lastException;
- private long lastWatermarkTime = 0L;
-
- private ReaderTask(
- Map readerToSsp,
- Coder checkpointMarkCoder,
- int capacity,
- long watermarkInterval,
- FnWithMetricsWrapper metricsWrapper) {
- this.readerToSsp = readerToSsp;
- this.checkpointMarkCoder = checkpointMarkCoder;
- this.readers = ImmutableList.copyOf(readerToSsp.keySet());
- this.watermarkInterval = watermarkInterval;
- this.available = new Semaphore(capacity);
- this.metricsWrapper = metricsWrapper;
-
- final Map> qs =
- new HashMap<>();
- readerToSsp.values().forEach(ssp -> qs.put(ssp, new LinkedBlockingQueue<>()));
- this.queues = ImmutableMap.copyOf(qs);
- }
-
- @Override
- public void run() {
- this.running = true;
-
- try {
- for (UnboundedReader reader : readers) {
- final boolean hasData = invoke(reader::start);
- if (hasData) {
- available.acquire();
- enqueueMessage(reader);
- }
- }
-
- while (running) {
- boolean elementAvailable = false;
- for (UnboundedReader reader : readers) {
- final boolean hasData = invoke(reader::advance);
- if (hasData) {
- while (!available.tryAcquire(
- 1,
- Math.max(lastWatermarkTime + watermarkInterval - System.currentTimeMillis(), 1),
- TimeUnit.MILLISECONDS)) {
- updateWatermark();
- }
- enqueueMessage(reader);
- elementAvailable = true;
- }
- }
-
- updateWatermark();
-
- if (!elementAvailable) {
- // TODO: make poll interval configurable
- Thread.sleep(50);
- }
- }
- } catch (Exception e) {
- lastException = e;
- running = false;
- } finally {
- readers.forEach(
- reader -> {
- try {
- reader.close();
- } catch (IOException e) {
- LOG.error("Reader task failed to close reader", e);
- }
- });
- }
-
- if (lastException != null) {
- // Force any pollers to wake up
- queues
- .values()
- .forEach(
- queue -> {
- queue.clear();
- queue.add(CHECK_LAST_EXCEPTION_ENVELOPE);
- });
- }
- }
-
- private X invoke(FnWithMetricsWrapper.SupplierWithException fn) throws Exception {
- if (metricsWrapper != null) {
- return metricsWrapper.wrap(fn, true);
- } else {
- return fn.get();
- }
- }
-
- private void updateWatermark() throws InterruptedException {
- final long time = System.currentTimeMillis();
- if (time - lastWatermarkTime > watermarkInterval) {
- for (UnboundedReader reader : readers) {
- final SystemStreamPartition ssp = readerToSsp.get(reader);
- final Instant currentWatermark =
- currentWatermarks.containsKey(ssp)
- ? currentWatermarks.get(ssp)
- : BoundedWindow.TIMESTAMP_MIN_VALUE;
- final Instant nextWatermark = reader.getWatermark();
- if (currentWatermark.isBefore(nextWatermark)) {
- currentWatermarks.put(ssp, nextWatermark);
- if (BoundedWindow.TIMESTAMP_MAX_VALUE.isAfter(nextWatermark)) {
- enqueueWatermark(reader);
- } else {
- // Max watermark has been reached for this reader.
- enqueueMaxWatermarkAndEndOfStream(reader);
- running = false;
- }
- }
- }
-
- lastWatermarkTime = time;
- }
- }
-
- private void enqueueWatermark(UnboundedReader reader) throws InterruptedException {
- final SystemStreamPartition ssp = readerToSsp.get(reader);
- final IncomingMessageEnvelope envelope =
- IncomingMessageEnvelope.buildWatermarkEnvelope(ssp, reader.getWatermark().getMillis());
-
- queues.get(ssp).put(envelope);
- }
-
- private void enqueueMessage(UnboundedReader reader) throws InterruptedException {
- @SuppressWarnings("unchecked")
- final T value = (T) reader.getCurrent();
- final Instant time = reader.getCurrentTimestamp();
- final SystemStreamPartition ssp = readerToSsp.get(reader);
- final WindowedValue windowedValue =
- WindowedValues.timestampedValueInGlobalWindow(value, time);
-
- final OpMessage opMessage = OpMessage.ofElement(windowedValue);
- final IncomingMessageEnvelope envelope =
- new IncomingMessageEnvelope(ssp, getOffset(reader), null, opMessage);
-
- queues.get(ssp).put(envelope);
- }
-
- // Send an max watermark message and an end of stream message to the corresponding ssp to
- // close windows and finish the task.
- private void enqueueMaxWatermarkAndEndOfStream(UnboundedReader reader) {
- final SystemStreamPartition ssp = readerToSsp.get(reader);
- // Send the max watermark to force completion of any open windows.
- final IncomingMessageEnvelope watermarkEnvelope =
- IncomingMessageEnvelope.buildWatermarkEnvelope(
- ssp, BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis());
- enqueueUninterruptibly(watermarkEnvelope);
-
- final IncomingMessageEnvelope endOfStreamEnvelope =
- IncomingMessageEnvelope.buildEndOfStreamEnvelope(ssp);
- enqueueUninterruptibly(endOfStreamEnvelope);
- }
-
- private void enqueueUninterruptibly(IncomingMessageEnvelope envelope) {
- final BlockingQueue queue =
- queues.get(envelope.getSystemStreamPartition());
- while (true) {
- try {
- queue.put(envelope);
- return;
- } catch (InterruptedException e) {
- // Some events require that we post an envelope to the queue even if the interrupt
- // flag was set (i.e. during a call to stop) to ensure that the consumer properly
- // shuts down. Consequently, if we receive an interrupt here we ignore it and retry
- // the put operation.
- }
- }
- }
-
- void stop() {
- running = false;
- }
-
- List getNextMessages(SystemStreamPartition ssp, long timeoutMillis)
- throws InterruptedException {
- if (lastException != null) {
- throw new RuntimeException(lastException);
- }
-
- final List envelopes = new ArrayList<>();
- final BlockingQueue queue = queues.get(ssp);
- final IncomingMessageEnvelope envelope = queue.poll(timeoutMillis, TimeUnit.MILLISECONDS);
-
- if (envelope != null) {
- envelopes.add(envelope);
- queue.drainTo(envelopes);
- }
-
- final int numElements =
- (int) envelopes.stream().filter(ev -> (ev.getMessage() instanceof OpMessage)).count();
- available.release(numElements);
-
- if (lastException != null) {
- throw new RuntimeException(lastException);
- }
-
- return envelopes;
- }
-
- private String getOffset(UnboundedReader reader) {
- try {
- final ByteArrayOutputStream baos = new ByteArrayOutputStream();
- @SuppressWarnings("unchecked")
- final CheckpointMarkT checkpointMark =
- (CheckpointMarkT) invoke(reader::getCheckpointMark);
- checkpointMarkCoder.encode(checkpointMark, baos);
- return Base64.getEncoder().encodeToString(baos.toByteArray());
- } catch (Exception e) {
- throw new RuntimeException(e);
- }
- }
- }
- }
-
- /**
- * A {@link SystemFactory} that produces a {@link UnboundedSourceSystem} for a particular {@link
- * UnboundedSource} registered in {@link Config}.
- */
- public static class Factory implements SystemFactory {
- @Override
- public SystemConsumer getConsumer(String systemName, Config config, MetricsRegistry registry) {
- final String streamPrefix = "systems." + systemName;
- final Config scopedConfig = config.subset(streamPrefix + ".", true);
- return new Consumer(
- getUnboundedSource(scopedConfig),
- getPipelineOptions(config),
- new SamzaMetricsContainer((MetricsRegistryMap) registry),
- scopedConfig.get("stepName"));
- }
-
- @Override
- public SystemProducer getProducer(String systemName, Config config, MetricsRegistry registry) {
- LOG.info("System {} does not have producer.", systemName);
- return null;
- }
-
- @Override
- public SystemAdmin getAdmin(String systemName, Config config) {
- final Config scopedConfig = config.subset("systems." + systemName + ".", true);
- return new Admin(
- getUnboundedSource(scopedConfig), getPipelineOptions(config));
- }
-
- private static
- UnboundedSource getUnboundedSource(Config config) {
- @SuppressWarnings("unchecked")
- final UnboundedSource source =
- Base64Serializer.deserializeUnchecked(config.get("source"), UnboundedSource.class);
- return source;
- }
-
- private static SamzaPipelineOptions getPipelineOptions(Config config) {
- return Base64Serializer.deserializeUnchecked(
- config.get("beamPipelineOptions"), SerializablePipelineOptions.class)
- .get()
- .as(SamzaPipelineOptions.class);
- }
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/adapter/package-info.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/adapter/package-info.java
deleted file mode 100644
index 582194440c9b..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/adapter/package-info.java
+++ /dev/null
@@ -1,20 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-/** Internal implementation of the Beam runner for Apache Samza. */
-package org.apache.beam.runners.samza.adapter;
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/container/BeamContainerRunner.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/container/BeamContainerRunner.java
deleted file mode 100644
index 0f4a1c7a6905..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/container/BeamContainerRunner.java
+++ /dev/null
@@ -1,89 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.container;
-
-import java.time.Duration;
-import org.apache.samza.application.SamzaApplication;
-import org.apache.samza.application.descriptors.ApplicationDescriptor;
-import org.apache.samza.application.descriptors.ApplicationDescriptorImpl;
-import org.apache.samza.application.descriptors.ApplicationDescriptorUtil;
-import org.apache.samza.config.Config;
-import org.apache.samza.config.ShellCommandConfig;
-import org.apache.samza.context.ExternalContext;
-import org.apache.samza.job.ApplicationStatus;
-import org.apache.samza.runtime.ApplicationRunner;
-import org.apache.samza.runtime.ContainerLaunchUtil;
-import org.apache.samza.util.SamzaUncaughtExceptionHandler;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/** Runs the beam Yarn container, using the static global job model. */
-@SuppressWarnings({
- "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
- "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
-public class BeamContainerRunner implements ApplicationRunner {
- private static final Logger LOG = LoggerFactory.getLogger(BeamContainerRunner.class);
-
- @SuppressWarnings("rawtypes")
- private final ApplicationDescriptorImpl appDesc;
-
- @SuppressWarnings("rawtypes")
- public BeamContainerRunner(SamzaApplication app, Config config) {
- this.appDesc = ApplicationDescriptorUtil.getAppDescriptor(app, config);
- }
-
- @Override
- public void run(ExternalContext externalContext) {
- Thread.setDefaultUncaughtExceptionHandler(
- new SamzaUncaughtExceptionHandler(
- () -> {
- LOG.info("Exiting process now.");
- System.exit(1);
- }));
-
- ContainerLaunchUtil.run(
- appDesc, System.getenv(ShellCommandConfig.ENV_CONTAINER_ID), ContainerCfgLoader.jobModel);
- }
-
- @Override
- public void kill() {
- // Do nothing. Yarn will kill the container.
- }
-
- @Override
- public ApplicationStatus status() {
- // The container is running during the life span of this object.
- return ApplicationStatus.Running;
- }
-
- @Override
- public void waitForFinish() {
- // Container run is synchronous
- // so calling waitForFinish() after run() should return immediately
- LOG.info("Container has stopped");
- }
-
- @Override
- public boolean waitForFinish(Duration timeout) {
- // Container run is synchronous
- // so calling waitForFinish() after run() should return immediately
- LOG.info("Container has stopped");
- return true;
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/container/BeamJobCoordinatorRunner.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/container/BeamJobCoordinatorRunner.java
deleted file mode 100644
index fb00a018fb29..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/container/BeamJobCoordinatorRunner.java
+++ /dev/null
@@ -1,78 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.container;
-
-import java.time.Duration;
-import org.apache.samza.application.SamzaApplication;
-import org.apache.samza.application.descriptors.ApplicationDescriptor;
-import org.apache.samza.clustermanager.JobCoordinatorLaunchUtil;
-import org.apache.samza.config.Config;
-import org.apache.samza.context.ExternalContext;
-import org.apache.samza.job.ApplicationStatus;
-import org.apache.samza.runtime.ApplicationRunner;
-
-/** Runs on Yarn AM, execute planning and launches JobCoordinator. */
-public class BeamJobCoordinatorRunner implements ApplicationRunner {
-
- @SuppressWarnings("rawtypes")
- private final SamzaApplication app;
-
- private final Config config;
-
- /**
- * Constructors a {@link BeamJobCoordinatorRunner} to run the {@code app} with the {@code config}.
- *
- * @param app application to run
- * @param config configuration for the application
- */
- @SuppressWarnings("rawtypes")
- public BeamJobCoordinatorRunner(
- SamzaApplication app, Config config) {
- this.app = app;
- this.config = config;
- }
-
- @Override
- public void run(ExternalContext externalContext) {
- JobCoordinatorLaunchUtil.run(app, config);
- }
-
- @Override
- public void kill() {
- throw new UnsupportedOperationException(
- "BeamJobCoordinatorRunner#kill should never be invoked.");
- }
-
- @Override
- public ApplicationStatus status() {
- throw new UnsupportedOperationException(
- "BeamJobCoordinatorRunner#status should never be invoked.");
- }
-
- @Override
- public void waitForFinish() {
- throw new UnsupportedOperationException(
- "BeamJobCoordinatorRunner#waitForFinish should never be invoked.");
- }
-
- @Override
- public boolean waitForFinish(Duration timeout) {
- throw new UnsupportedOperationException(
- "BeamJobCoordinatorRunner#waitForFinish should never be invoked.");
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/container/ContainerCfgLoader.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/container/ContainerCfgLoader.java
deleted file mode 100644
index 9437aea56561..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/container/ContainerCfgLoader.java
+++ /dev/null
@@ -1,62 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.container;
-
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Random;
-import org.apache.samza.config.Config;
-import org.apache.samza.config.ConfigLoader;
-import org.apache.samza.config.MapConfig;
-import org.apache.samza.config.ShellCommandConfig;
-import org.apache.samza.container.SamzaContainer;
-import org.apache.samza.job.model.JobModel;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/** Loader for the Beam yarn container to load job model. */
-@SuppressWarnings({
- "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
-public class ContainerCfgLoader implements ConfigLoader {
- private static final Logger LOG = LoggerFactory.getLogger(ContainerCfgLoader.class);
-
- private static final Object LOCK = new Object();
- static volatile JobModel jobModel;
- private static final Random RANDOM = new Random();
-
- @Override
- public Config getConfig() {
- if (jobModel == null) {
- synchronized (LOCK) {
- if (jobModel == null) {
- final String containerId = System.getenv(ShellCommandConfig.ENV_CONTAINER_ID);
- LOG.info("Got container ID: {}", containerId);
- final String coordinatorUrl = System.getenv(ShellCommandConfig.ENV_COORDINATOR_URL);
- LOG.info("Got coordinator URL: {}", coordinatorUrl);
- final int delay = RANDOM.nextInt(SamzaContainer.DEFAULT_READ_JOBMODEL_DELAY_MS()) + 1;
- jobModel = SamzaContainer.readJobModel(coordinatorUrl, delay);
- }
- }
- }
-
- final Map config = new HashMap<>(jobModel.getConfig());
- config.put("app.runner.class", BeamContainerRunner.class.getName());
- return new MapConfig(config);
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/container/ContainerCfgLoaderFactory.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/container/ContainerCfgLoaderFactory.java
deleted file mode 100644
index d3b090d6e20a..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/container/ContainerCfgLoaderFactory.java
+++ /dev/null
@@ -1,30 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.container;
-
-import org.apache.samza.config.Config;
-import org.apache.samza.config.ConfigLoader;
-import org.apache.samza.config.ConfigLoaderFactory;
-
-/** Factory for the Beam yarn container to get loader to load job model. */
-public class ContainerCfgLoaderFactory implements ConfigLoaderFactory {
- @Override
- public ConfigLoader getLoader(Config config) {
- return new ContainerCfgLoader();
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/container/package-info.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/container/package-info.java
deleted file mode 100644
index 58a09e6023de..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/container/package-info.java
+++ /dev/null
@@ -1,20 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-/** Internal implementation of the Beam runner for Apache Samza. */
-package org.apache.beam.runners.samza.container;
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java
deleted file mode 100644
index 7bec91abb34d..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java
+++ /dev/null
@@ -1,108 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.metrics;
-
-import org.apache.beam.runners.core.DoFnRunner;
-import org.apache.beam.sdk.state.TimeDomain;
-import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.values.CausedByDrain;
-import org.apache.beam.sdk.values.WindowedValue;
-import org.joda.time.Instant;
-
-/**
- * {@link DoFnRunner} wrapper with metrics. The class uses {@link SamzaMetricsContainer} to keep
- * BEAM metrics results and update Samza metrics.
- */
-public class DoFnRunnerWithMetrics implements DoFnRunner {
- private final DoFnRunner underlying;
- private final FnWithMetricsWrapper metricsWrapper;
-
- private DoFnRunnerWithMetrics(
- DoFnRunner underlying, SamzaMetricsContainer metricsContainer, String stepName) {
- this.underlying = underlying;
- this.metricsWrapper = new FnWithMetricsWrapper(metricsContainer, stepName);
- }
-
- public static DoFnRunner wrap(
- DoFnRunner doFnRunner, SamzaMetricsContainer metricsContainer, String stepName) {
- return new DoFnRunnerWithMetrics<>(doFnRunner, metricsContainer, stepName);
- }
-
- @Override
- public void startBundle() {
- withMetrics(underlying::startBundle, false);
- }
-
- @Override
- public void processElement(WindowedValue elem) {
- withMetrics(() -> underlying.processElement(elem), false);
- }
-
- @Override
- public void onTimer(
- String timerId,
- String timerFamilyId,
- KeyT key,
- BoundedWindow window,
- Instant timestamp,
- Instant outputTimestamp,
- TimeDomain timeDomain,
- CausedByDrain causedByDrain) {
- withMetrics(
- () ->
- underlying.onTimer(
- timerId,
- timerFamilyId,
- key,
- window,
- timestamp,
- outputTimestamp,
- timeDomain,
- causedByDrain),
- false);
- }
-
- @Override
- public void finishBundle() {
- withMetrics(underlying::finishBundle, true);
- }
-
- @Override
- public void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {
- underlying.onWindowExpiration(window, timestamp, key);
- }
-
- @Override
- public DoFn getFn() {
- return underlying.getFn();
- }
-
- private void withMetrics(Runnable runnable, boolean shouldUpdateMetrics) {
- try {
- metricsWrapper.wrap(
- () -> {
- runnable.run();
- return (Void) null;
- },
- shouldUpdateMetrics);
- } catch (Exception e) {
- throw new RuntimeException(e);
- }
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/FnWithMetricsWrapper.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/FnWithMetricsWrapper.java
deleted file mode 100644
index 8052e0d66146..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/FnWithMetricsWrapper.java
+++ /dev/null
@@ -1,50 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.metrics;
-
-import java.io.Closeable;
-import org.apache.beam.sdk.metrics.MetricsEnvironment;
-
-/** This class wraps a {@link java.util.function.Supplier} function call with BEAM metrics. */
-public class FnWithMetricsWrapper {
-
- /** Interface for functions to be wrapped with metrics. */
- public interface SupplierWithException {
- T get() throws Exception;
- }
-
- private final SamzaMetricsContainer metricsContainer;
- private final String stepName;
-
- public FnWithMetricsWrapper(SamzaMetricsContainer metricsContainer, String stepName) {
- this.metricsContainer = metricsContainer;
- this.stepName = stepName;
- }
-
- public T wrap(SupplierWithException fn, boolean shouldUpdateMetrics) throws Exception {
- try (Closeable closeable =
- MetricsEnvironment.scopedMetricsContainer(metricsContainer.getContainer(stepName))) {
- T result = fn.get();
- // Skip updating metrics if not necessary to improve performance
- if (shouldUpdateMetrics) {
- metricsContainer.updateMetrics(stepName);
- }
- return result;
- }
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaGBKMetricOp.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaGBKMetricOp.java
deleted file mode 100644
index 9b6375171c23..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaGBKMetricOp.java
+++ /dev/null
@@ -1,194 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.metrics;
-
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
-import java.math.BigInteger;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.concurrent.ConcurrentHashMap;
-import org.apache.beam.runners.samza.runtime.KeyedTimerData;
-import org.apache.beam.runners.samza.runtime.Op;
-import org.apache.beam.runners.samza.runtime.OpEmitter;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.values.WindowedValue;
-import org.apache.samza.config.Config;
-import org.apache.samza.context.Context;
-import org.apache.samza.operators.Scheduler;
-import org.joda.time.Instant;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * SamzaGBKMetricOp is a {@link Op} that emits & maintains default metrics for input or output
- * PCollection for GroupByKey.
- *
- * For Input PCollection: It emits the input throughput and maintains avg input time for input
- * PCollection per windowId.
- *
- *
For Output PCollection: It emits the output throughput and maintains avg output time for
- * output PCollection per windowId. It is also responsible for emitting latency metric per windowId
- * once the watermark passes the end of window timestamp.
- *
- *
Assumes that {@code SamzaGBKMetricOp#processWatermark(Instant, OpEmitter)} is exclusive of
- * {@code SamzaGBKMetricOp#processElement(Instant, OpEmitter)}. Specifically, the processWatermark
- * method assumes that no calls to processElement will be made during its execution, and vice versa.
- *
- * @param The type of the elements in the input PCollection.
- */
-class SamzaGBKMetricOp implements Op {
- private static final Logger LOG = LoggerFactory.getLogger(SamzaGBKMetricOp.class);
- // Unique name of the PTransform this MetricOp is associated with
- private final String transformFullName;
- private final SamzaTransformMetricRegistry samzaTransformMetricRegistry;
- // Type of the processing operation
- private final SamzaMetricOpFactory.OpType opType;
-
- private final String pValue;
- // Counters for keeping sum of arrival time and count of elements per windowId
- @SuppressFBWarnings("SE_BAD_FIELD")
- private final ConcurrentHashMap sumOfTimestampsPerWindowId;
-
- @SuppressFBWarnings("SE_BAD_FIELD")
- private final ConcurrentHashMap sumOfCountPerWindowId;
- // Name of the task, for logging purpose
- private transient String task;
-
- @Override
- @SuppressWarnings({"rawtypes", "unchecked"})
- public void open(
- Config config,
- Context context,
- Scheduler> timerRegistry,
- OpEmitter emitter) {
- // for logging / debugging purposes
- this.task = context.getTaskContext().getTaskModel().getTaskName().getTaskName();
- // Register the transform with SamzaTransformMetricRegistry
- samzaTransformMetricRegistry.register(transformFullName, pValue, context);
- }
-
- // Some fields are initialized in open() method, which is called after the constructor.
- @SuppressWarnings("initialization.fields.uninitialized")
- public SamzaGBKMetricOp(
- String pValue,
- String transformFullName,
- SamzaMetricOpFactory.OpType opType,
- SamzaTransformMetricRegistry samzaTransformMetricRegistry) {
- this.pValue = pValue;
- this.transformFullName = transformFullName;
- this.opType = opType;
- this.samzaTransformMetricRegistry = samzaTransformMetricRegistry;
- this.sumOfTimestampsPerWindowId = new ConcurrentHashMap<>();
- this.sumOfCountPerWindowId = new ConcurrentHashMap<>();
- }
-
- @Override
- public void processElement(WindowedValue inputElement, OpEmitter emitter) {
- // one element can belong to multiple windows
- for (BoundedWindow windowId : inputElement.getWindows()) {
- // Atomic updates to counts
- sumOfCountPerWindowId.compute(
- windowId,
- (key, value) -> {
- value = value == null ? Long.valueOf(0) : value;
- return ++value;
- });
- // Atomic updates to sum of arrival timestamps
- sumOfTimestampsPerWindowId.compute(
- windowId,
- (key, value) -> {
- value = value == null ? BigInteger.ZERO : value;
- return value.add(BigInteger.valueOf(System.nanoTime()));
- });
- }
-
- switch (opType) {
- case INPUT:
- samzaTransformMetricRegistry
- .getTransformMetrics()
- .getTransformInputThroughput(transformFullName)
- .inc();
- break;
- case OUTPUT:
- samzaTransformMetricRegistry
- .getTransformMetrics()
- .getTransformOutputThroughput(transformFullName)
- .inc();
- break;
- }
- emitter.emitElement(inputElement);
- }
-
- @Override
- public void processWatermark(Instant watermark, OpEmitter emitter) {
- final List closedWindows = new ArrayList<>();
- sumOfCountPerWindowId.keySet().stream()
- .filter(windowId -> watermark.isAfter(windowId.maxTimestamp())) // window is closed
- .forEach(
- windowId -> {
- // In case if BigInteger overflows for long we only retain the last 64 bits of the sum
- long sumOfTimestamps =
- sumOfTimestampsPerWindowId.get(windowId) != null
- ? sumOfTimestampsPerWindowId.get(windowId).longValue()
- : 0L;
- long count = sumOfCountPerWindowId.get(windowId);
- closedWindows.add(windowId);
-
- if (LOG.isDebugEnabled()) {
- LOG.debug(
- "Processing {} Watermark for Transform: {}, WindowId:{}, count: {}, sumOfTimestamps: {}, task: {}",
- opType,
- transformFullName,
- windowId,
- count,
- sumOfTimestamps,
- task);
- }
-
- // if the window is closed and there is some data
- if (sumOfTimestamps > 0 && count > 0) {
- switch (opType) {
- case INPUT:
- // Update the arrival time for the window
- samzaTransformMetricRegistry.updateArrivalTimeMap(
- transformFullName, windowId, Math.floorDiv(sumOfTimestamps, count));
- break;
- case OUTPUT:
- // Compute the latency if there is some data for the window
- samzaTransformMetricRegistry.emitLatencyMetric(
- transformFullName, windowId, Math.floorDiv(sumOfTimestamps, count), task);
- break;
- }
- }
- });
-
- // remove the closed windows
- sumOfCountPerWindowId.keySet().removeAll(closedWindows);
- sumOfTimestampsPerWindowId.keySet().removeAll(closedWindows);
-
- // Update the watermark progress for the transform output
- if (opType == SamzaMetricOpFactory.OpType.OUTPUT) {
- samzaTransformMetricRegistry
- .getTransformMetrics()
- .getTransformWatermarkProgress(transformFullName)
- .set(watermark.getMillis());
- }
-
- emitter.emitWatermark(watermark);
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaMetricOp.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaMetricOp.java
deleted file mode 100644
index 0f5334546c7c..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaMetricOp.java
+++ /dev/null
@@ -1,173 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.metrics;
-
-import java.math.BigInteger;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.atomic.AtomicLong;
-import java.util.concurrent.atomic.AtomicReference;
-import org.apache.beam.runners.samza.runtime.KeyedTimerData;
-import org.apache.beam.runners.samza.runtime.Op;
-import org.apache.beam.runners.samza.runtime.OpEmitter;
-import org.apache.beam.runners.samza.util.PipelineJsonRenderer;
-import org.apache.beam.sdk.values.WindowedValue;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
-import org.apache.samza.config.Config;
-import org.apache.samza.context.Context;
-import org.apache.samza.operators.Scheduler;
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.joda.time.Instant;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * SamzaMetricOp is a metric Op that emits & maintains default transform metrics for inputs &
- * outputs PCollection to the non data-shuffle transform. It emits the output throughput and
- * maintains avg arrival time for input & output PCollection per watermark.
- *
- * Assumes that {@code SamzaMetricOp#processWatermark(Instant, OpEmitter)} is exclusive of {@code
- * SamzaMetricOp#processElement(Instant, OpEmitter)}. Specifically, the processWatermark method
- * assumes that no calls to processElement will be made during its execution, and vice versa.
- *
- * @param The type of the elements in the output PCollection.
- */
-class SamzaMetricOp implements Op {
- // Unique name of the PTransform this MetricOp is associated with
- private final String transformFullName;
- private final SamzaTransformMetricRegistry samzaTransformMetricRegistry;
- // Name or identifier of the PCollection which PTransform is processing
- private final String pValue;
- // Counters for output throughput
- private final AtomicLong count;
- private final AtomicReference sumOfTimestamps;
- // Type of the PTransform input or output
- private final SamzaMetricOpFactory.OpType opType;
- // List of input PValue(s) for all PCollections processing the PTransform
- private transient List transformInputs;
- // List of output PValue(s) for all PCollections processing the PTransform
- private transient List transformOutputs;
- // Name of the task, for logging purpose
- private transient String task;
-
- private static final Logger LOG = LoggerFactory.getLogger(SamzaMetricOp.class);
-
- // Some fields are initialized in open() method, which is called after the constructor.
- @SuppressWarnings("initialization.fields.uninitialized")
- public SamzaMetricOp(
- @NonNull String pValue,
- @NonNull String transformFullName,
- SamzaMetricOpFactory.OpType opType,
- @NonNull SamzaTransformMetricRegistry samzaTransformMetricRegistry) {
- this.transformFullName = transformFullName;
- this.samzaTransformMetricRegistry = samzaTransformMetricRegistry;
- this.pValue = pValue;
- this.opType = opType;
- this.count = new AtomicLong(0L);
- this.sumOfTimestamps = new AtomicReference<>(BigInteger.ZERO);
- }
-
- @Override
- @SuppressWarnings({"rawtypes", "unchecked"})
- public void open(
- Config config,
- Context context,
- Scheduler> timerRegistry,
- OpEmitter emitter) {
- final Map.Entry, List> transformInputOutput =
- PipelineJsonRenderer.getTransformIOMap(config).get(transformFullName);
- this.transformInputs =
- transformInputOutput != null ? transformInputOutput.getKey() : new ArrayList();
- this.transformOutputs =
- transformInputOutput != null ? transformInputOutput.getValue() : new ArrayList();
- // for logging / debugging purposes
- this.task = context.getTaskContext().getTaskModel().getTaskName().getTaskName();
- // Register the transform with SamzaTransformMetricRegistry
- samzaTransformMetricRegistry.register(transformFullName, pValue, context);
- }
-
- @Override
- public void processElement(WindowedValue inputElement, OpEmitter emitter) {
- // update counters for timestamps
- count.incrementAndGet();
- sumOfTimestamps.updateAndGet(sum -> sum.add(BigInteger.valueOf(System.nanoTime())));
- switch (opType) {
- case INPUT:
- samzaTransformMetricRegistry
- .getTransformMetrics()
- .getTransformInputThroughput(transformFullName)
- .inc();
- break;
- case OUTPUT:
- samzaTransformMetricRegistry
- .getTransformMetrics()
- .getTransformOutputThroughput(transformFullName)
- .inc();
- break;
- }
- emitter.emitElement(inputElement);
- }
-
- @Override
- @SuppressWarnings({"CompareToZero"})
- public void processWatermark(Instant watermark, OpEmitter emitter) {
- if (LOG.isDebugEnabled()) {
- LOG.debug(
- "Processing Output Watermark for Transform: {} Count: {} SumOfTimestamps: {} for Watermark: {} for Task: {}",
- transformFullName,
- count.get(),
- sumOfTimestamps.get().longValue(),
- watermark.getMillis(),
- task);
- }
-
- // if there is no input data then counters will be zero and only watermark will progress
- if (count.get() > 0) {
- // if BigInt.longValue is out of range for long then only the low-order 64 bits are retained
- long avg = Math.floorDiv(sumOfTimestamps.get().longValue(), count.get());
- // Update MetricOp Registry with avg arrival for the pValue
- samzaTransformMetricRegistry.updateArrivalTimeMap(
- transformFullName, pValue, watermark.getMillis(), avg);
- if (opType == SamzaMetricOpFactory.OpType.OUTPUT) {
- // compute & emit the latency metric if the opType is OUTPUT
- samzaTransformMetricRegistry.emitLatencyMetric(
- transformFullName, transformInputs, transformOutputs, watermark.getMillis(), task);
- }
- }
-
- if (opType == SamzaMetricOpFactory.OpType.OUTPUT) {
- // update output watermark progress metric
- samzaTransformMetricRegistry
- .getTransformMetrics()
- .getTransformWatermarkProgress(transformFullName)
- .set(watermark.getMillis());
- }
-
- // reset all counters
- count.set(0L);
- this.sumOfTimestamps.set(BigInteger.ZERO);
- emitter.emitWatermark(watermark);
- }
-
- @VisibleForTesting
- void init(List transformInputs, List transformOutputs) {
- this.transformInputs = transformInputs;
- this.transformOutputs = transformOutputs;
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaMetricOpFactory.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaMetricOpFactory.java
deleted file mode 100644
index a4112a510459..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaMetricOpFactory.java
+++ /dev/null
@@ -1,73 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.metrics;
-
-import org.apache.beam.runners.samza.runtime.Op;
-import org.apache.beam.sdk.util.construction.PTransformTranslation;
-import org.checkerframework.checker.nullness.qual.NonNull;
-
-/**
- * Factory class to create {@link Op} for default transform metric computation.
- *
- * Each metric Op computes and emits default throughput, latency & watermark progress metric per
- * transform for Beam Samza Runner. A metric Op can be either attached to Input PCollection or
- * Output PCollection of a PTransform.
- *
- *
Each concrete metric OP is responsible for following metrics computation: 1. Throughput: Emit
- * the number of elements processed in the PCollection 2. Watermark Progress: Emit the output
- * watermark progress of the PCollection 3. Latency: Maintain the avg arrival time per watermark
- * across elements it processes, compute & emit the latency
- */
-@SuppressWarnings({
- "rawtypes", // TODO(https://issues.apache.org/jira/browse/BEAM-10556)
- "nullness"
-}) // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
-public class SamzaMetricOpFactory {
- public enum OpType {
- INPUT,
- OUTPUT
- }
-
- /**
- * Create a {@link Op} for default transform metric computation.
- *
- * @param urn URN of the PCollection metric Op is processing
- * @param pValue name of the PCollection metric Op is processing
- * @param transformName name of the PTransform for which metric Op is created
- * @param opType type of the metric
- * @param samzaTransformMetricRegistry metric registry
- * @param type of the message
- * @return a {@link Op} for default transform metric computation
- */
- public static @NonNull Op createMetricOp(
- @NonNull String urn,
- @NonNull String pValue,
- @NonNull String transformName,
- @NonNull OpType opType,
- @NonNull SamzaTransformMetricRegistry samzaTransformMetricRegistry) {
- if (isDataShuffleTransform(urn)) {
- return new SamzaGBKMetricOp<>(pValue, transformName, opType, samzaTransformMetricRegistry);
- }
- return new SamzaMetricOp<>(pValue, transformName, opType, samzaTransformMetricRegistry);
- }
-
- private static boolean isDataShuffleTransform(String urn) {
- return urn.equals(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN)
- || urn.equals(PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN);
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaMetricsContainer.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaMetricsContainer.java
deleted file mode 100644
index 1679b748b10b..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaMetricsContainer.java
+++ /dev/null
@@ -1,120 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.metrics;
-
-import static org.apache.beam.runners.core.metrics.MetricsContainerStepMap.asAttemptedOnlyMetricResults;
-
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.function.Consumer;
-import org.apache.beam.runners.core.metrics.MetricsContainerStepMap;
-import org.apache.beam.sdk.metrics.GaugeResult;
-import org.apache.beam.sdk.metrics.MetricQueryResults;
-import org.apache.beam.sdk.metrics.MetricResult;
-import org.apache.beam.sdk.metrics.MetricResults;
-import org.apache.beam.sdk.metrics.MetricsContainer;
-import org.apache.beam.sdk.metrics.MetricsFilter;
-import org.apache.samza.metrics.Counter;
-import org.apache.samza.metrics.Gauge;
-import org.apache.samza.metrics.Metric;
-import org.apache.samza.metrics.MetricsRegistryMap;
-
-/**
- * This class holds the {@link MetricsContainer}s for BEAM metrics, and update the results to Samza
- * metrics.
- */
-@SuppressWarnings({
- "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
-public class SamzaMetricsContainer {
- private static final String BEAM_METRICS_GROUP = "BeamMetrics";
-
- private final MetricsContainerStepMap metricsContainers = new MetricsContainerStepMap();
- private final MetricsRegistryMap metricsRegistry;
-
- public SamzaMetricsContainer(MetricsRegistryMap metricsRegistry) {
- this.metricsRegistry = metricsRegistry;
- this.metricsRegistry.metrics().put(BEAM_METRICS_GROUP, new ConcurrentHashMap<>());
- }
-
- public MetricsContainer getContainer(String stepName) {
- return this.metricsContainers.getContainer(stepName);
- }
-
- public MetricsContainerStepMap getContainers() {
- return this.metricsContainers;
- }
-
- public void updateMetrics(String stepName) {
- assert metricsRegistry != null;
-
- final MetricResults metricResults = asAttemptedOnlyMetricResults(metricsContainers);
- final MetricQueryResults results =
- metricResults.queryMetrics(MetricsFilter.builder().addStep(stepName).build());
-
- final CounterUpdater updateCounter = new CounterUpdater();
- results.getCounters().forEach(updateCounter);
-
- final GaugeUpdater updateGauge = new GaugeUpdater();
- results.getGauges().forEach(updateGauge);
-
- // TODO(https://github.com/apache/beam/issues/21043): add distribution metrics to Samza
- }
-
- public void updateExecutableStageBundleMetric(String metricName, long time) {
- @SuppressWarnings("unchecked")
- Gauge gauge = (Gauge) getSamzaMetricFor(metricName);
- if (gauge == null) {
- gauge = metricsRegistry.newGauge(BEAM_METRICS_GROUP, metricName, 0L);
- }
- gauge.set(time);
- }
-
- private class CounterUpdater implements Consumer> {
- @Override
- public void accept(MetricResult metricResult) {
- final String metricName = getMetricName(metricResult);
- Counter counter = (Counter) getSamzaMetricFor(metricName);
- if (counter == null) {
- counter = metricsRegistry.newCounter(BEAM_METRICS_GROUP, metricName);
- }
- counter.dec(counter.getCount());
- counter.inc(metricResult.getAttempted());
- }
- }
-
- private class GaugeUpdater implements Consumer> {
- @Override
- public void accept(MetricResult metricResult) {
- final String metricName = getMetricName(metricResult);
- @SuppressWarnings("unchecked")
- Gauge gauge = (Gauge) getSamzaMetricFor(metricName);
- if (gauge == null) {
- gauge = metricsRegistry.newGauge(BEAM_METRICS_GROUP, metricName, 0L);
- }
- gauge.set(metricResult.getAttempted().getValue());
- }
- }
-
- private Metric getSamzaMetricFor(String metricName) {
- return metricsRegistry.getGroup(BEAM_METRICS_GROUP).get(metricName);
- }
-
- private static String getMetricName(MetricResult metricResult) {
- return metricResult.getKey().toString();
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaTransformMetricRegistry.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaTransformMetricRegistry.java
deleted file mode 100644
index add207752f06..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaTransformMetricRegistry.java
+++ /dev/null
@@ -1,208 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.metrics;
-
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
-import java.io.Serializable;
-import java.util.Collections;
-import java.util.List;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.stream.Collectors;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
-import org.apache.samza.context.Context;
-import org.checkerframework.checker.nullness.qual.Nullable;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * SamzaTransformMetricRegistry is a registry that maintains the metrics for each transform. It
- * maintains the average arrival time for each PCollection for a primitive transform.
- *
- * For a non-data shuffling primitive transform, the average arrival time is calculated per
- * watermark, per PCollection {@link org.apache.beam.sdk.values.PValue} and updated in
- * avgArrivalTimeMap
- */
-public class SamzaTransformMetricRegistry implements Serializable {
- private static final Logger LOG = LoggerFactory.getLogger(SamzaTransformMetricRegistry.class);
-
- // TransformName -> PValue for pCollection -> Map
- private final ConcurrentHashMap>>
- avgArrivalTimeMap;
- // TransformName -> Map
- @SuppressFBWarnings("SE_BAD_FIELD")
- private final ConcurrentHashMap>
- avgArrivalTimeMapForGbk;
-
- // Per Transform Metrics for each primitive transform
- private final SamzaTransformMetrics transformMetrics;
-
- public SamzaTransformMetricRegistry() {
- this.avgArrivalTimeMap = new ConcurrentHashMap<>();
- this.avgArrivalTimeMapForGbk = new ConcurrentHashMap<>();
- this.transformMetrics = new SamzaTransformMetrics();
- }
-
- @VisibleForTesting
- SamzaTransformMetricRegistry(SamzaTransformMetrics samzaTransformMetrics) {
- this.transformMetrics = samzaTransformMetrics;
- this.avgArrivalTimeMap = new ConcurrentHashMap<>();
- this.avgArrivalTimeMapForGbk = new ConcurrentHashMap<>();
- }
-
- public void register(String transformFullName, String pValue, Context ctx) {
- transformMetrics.register(transformFullName, ctx);
- // initialize the map for the transform
- avgArrivalTimeMap.putIfAbsent(transformFullName, new ConcurrentHashMap<>());
- avgArrivalTimeMap.get(transformFullName).putIfAbsent(pValue, new ConcurrentHashMap<>());
- avgArrivalTimeMapForGbk.putIfAbsent(transformFullName, new ConcurrentHashMap<>());
- }
-
- public SamzaTransformMetrics getTransformMetrics() {
- return transformMetrics;
- }
-
- public void updateArrivalTimeMap(String transformName, String pValue, long watermark, long avg) {
- if (avgArrivalTimeMap.get(transformName) != null
- && avgArrivalTimeMap.get(transformName).get(pValue) != null) {
- ConcurrentHashMap avgArrivalTimeMapForPValue =
- avgArrivalTimeMap.get(transformName).get(pValue);
- // update the average arrival time for the latest watermark
- avgArrivalTimeMapForPValue.put(watermark, avg);
- // remove any stale entries which are lesser than the watermark
- avgArrivalTimeMapForPValue.entrySet().removeIf(entry -> entry.getKey() < watermark);
- }
- }
-
- public void updateArrivalTimeMap(String transformName, BoundedWindow windowId, long avg) {
- ConcurrentHashMap avgArrivalTimeMapForTransform =
- avgArrivalTimeMapForGbk.get(transformName);
- if (avgArrivalTimeMapForTransform != null) {
- avgArrivalTimeMapForTransform.put(windowId, avg);
- }
- }
-
- @SuppressWarnings("nullness")
- public void emitLatencyMetric(
- String transformName, BoundedWindow windowId, long avgArrivalEndTime, String taskName) {
- Long avgArrivalStartTime =
- avgArrivalTimeMapForGbk.get(transformName) != null
- ? avgArrivalTimeMapForGbk.get(transformName).remove(windowId)
- : null;
-
- if (avgArrivalStartTime == null || avgArrivalStartTime == 0 || avgArrivalEndTime == 0) {
- LOG.debug(
- "Failure to Emit Metric for Transform: {}, Start-Time: {} or End-Time: {} found is 0/null for windowId: {}, task: {}",
- transformName,
- avgArrivalStartTime,
- avgArrivalEndTime,
- windowId,
- taskName);
- return;
- }
-
- if (LOG.isDebugEnabled()) {
- LOG.debug(
- "Success Emit Metric for Transform: {}, window: {} for task: {}",
- transformName,
- windowId,
- taskName);
- }
- transformMetrics
- .getTransformLatencyMetric(transformName)
- .update(avgArrivalEndTime - avgArrivalStartTime);
-
- transformMetrics
- .getTransformCacheSize(transformName)
- .set((long) avgArrivalTimeMapForGbk.get(transformName).size());
- }
-
- // Checker framework bug: https://github.com/typetools/checker-framework/issues/979
- @SuppressWarnings("return")
- public void emitLatencyMetric(
- String transformName,
- List inputs,
- List outputs,
- Long watermark,
- String taskName) {
- final ConcurrentHashMap> avgArrivalTimeMapForTransform =
- avgArrivalTimeMap.get(transformName);
-
- if (avgArrivalTimeMapForTransform == null || inputs.isEmpty() || outputs.isEmpty()) {
- return;
- }
-
- // get the avg arrival times for all the input PValues
- final List inputPValuesAvgArrivalTimes =
- inputs.stream()
- .map(avgArrivalTimeMapForTransform::get)
- .map(map -> map == null ? null : map.remove(watermark))
- .filter(avgArrivalTime -> avgArrivalTime != null)
- .collect(Collectors.toList());
-
- // get the avg arrival times for all the output PValues
- final List outputPValuesAvgArrivalTimes =
- outputs.stream()
- .map(avgArrivalTimeMapForTransform::get)
- .map(map -> map == null ? null : map.remove(watermark))
- .filter(avgArrivalTime -> avgArrivalTime != null)
- .collect(Collectors.toList());
-
- if (inputPValuesAvgArrivalTimes.isEmpty() || outputPValuesAvgArrivalTimes.isEmpty()) {
- LOG.debug(
- "Failure to Emit Metric for Transform: {} inputArrivalTime: {} or outputArrivalTime: {} not found for Watermark: {} Task: {}",
- transformName,
- inputPValuesAvgArrivalTimes,
- inputPValuesAvgArrivalTimes,
- watermark,
- taskName);
- return;
- }
-
- final long startTime = Collections.min(inputPValuesAvgArrivalTimes);
- final long endTime = Collections.max(outputPValuesAvgArrivalTimes);
- final long latency = endTime - startTime;
- transformMetrics.getTransformLatencyMetric(transformName).update(latency);
-
- transformMetrics
- .getTransformCacheSize(transformName)
- .set(
- avgArrivalTimeMapForTransform.values().stream()
- .mapToLong(ConcurrentHashMap::size)
- .sum());
-
- LOG.debug(
- "Success Emit Metric Transform: {} for watermark: {} for task: {}",
- transformName,
- watermark,
- taskName);
- }
-
- @VisibleForTesting
- @Nullable
- ConcurrentHashMap> getAverageArrivalTimeMap(
- String transformName) {
- return avgArrivalTimeMap.get(transformName);
- }
-
- @VisibleForTesting
- @Nullable
- ConcurrentHashMap getAverageArrivalTimeMapForGBK(String transformName) {
- return avgArrivalTimeMapForGbk.get(transformName);
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaTransformMetrics.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaTransformMetrics.java
deleted file mode 100644
index 229b6da4e7c0..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/SamzaTransformMetrics.java
+++ /dev/null
@@ -1,142 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.metrics;
-
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
-import java.io.Serializable;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-import org.apache.samza.context.Context;
-import org.apache.samza.metrics.Counter;
-import org.apache.samza.metrics.Gauge;
-import org.apache.samza.metrics.MetricsRegistry;
-import org.apache.samza.metrics.SlidingTimeWindowReservoir;
-import org.apache.samza.metrics.Timer;
-
-/**
- * Metrics like throughput, latency and watermark progress for each Beam transform for Samza Runner.
- */
-@SuppressWarnings("return")
-public class SamzaTransformMetrics implements Serializable {
- private static final String ENABLE_TASK_METRICS = "runner.samza.transform.enable.task.metrics";
-
- private static final int DEFAULT_LOOKBACK_TIMER_WINDOW_SIZE_MS = 180000;
- private static final String GROUP = "SamzaBeamTransformMetrics";
- private static final String TRANSFORM_LATENCY_METRIC = "handle-message-ns";
- private static final String TRANSFORM_WATERMARK_PROGRESS = "output-watermark-ms";
- private static final String TRANSFORM_IP_THROUGHPUT = "num-input-messages";
- private static final String TRANSFORM_OP_THROUGHPUT = "num-output-messages";
-
- private static final String TRANSFORM_ARRIVAL_TIME_CACHE_SIZE = "in-mem-cache-size";
-
- // Transform name to metric maps
- @SuppressFBWarnings("SE_BAD_FIELD")
- private final Map transformLatency;
-
- @SuppressFBWarnings("SE_BAD_FIELD")
- private final Map> transformWatermarkProgress;
-
- @SuppressFBWarnings("SE_BAD_FIELD")
- private final Map transformInputThroughput;
-
- @SuppressFBWarnings("SE_BAD_FIELD")
- private final Map transformOutputThroughPut;
-
- @SuppressFBWarnings("SE_BAD_FIELD")
- private final Map> transformCacheSize;
-
- public SamzaTransformMetrics() {
- this.transformLatency = new ConcurrentHashMap<>();
- this.transformOutputThroughPut = new ConcurrentHashMap<>();
- this.transformWatermarkProgress = new ConcurrentHashMap<>();
- this.transformInputThroughput = new ConcurrentHashMap<>();
- this.transformCacheSize = new ConcurrentHashMap<>();
- }
-
- public void register(String transformName, Context ctx) {
- // Output Watermark metric per transform will always be per transform, per task, since per
- // container output watermark is not useful for debugging
- transformWatermarkProgress.putIfAbsent(
- transformName,
- ctx.getTaskContext()
- .getTaskMetricsRegistry()
- .newGauge(
- GROUP, getMetricNameWithPrefix(TRANSFORM_WATERMARK_PROGRESS, transformName), 0L));
-
- // Latency, throughput metrics can be per container (default) or per task
- final boolean enablePerTaskMetrics =
- ctx.getJobContext().getConfig().getBoolean(ENABLE_TASK_METRICS, false);
- final MetricsRegistry metricsRegistry =
- enablePerTaskMetrics
- ? ctx.getTaskContext().getTaskMetricsRegistry()
- : ctx.getContainerContext().getContainerMetricsRegistry();
- transformLatency.putIfAbsent(
- transformName,
- metricsRegistry.newTimer(GROUP, getTimerWithCustomizedLookBackWindow(transformName)));
- transformOutputThroughPut.putIfAbsent(
- transformName,
- metricsRegistry.newCounter(
- GROUP, getMetricNameWithPrefix(TRANSFORM_OP_THROUGHPUT, transformName)));
- transformInputThroughput.putIfAbsent(
- transformName,
- metricsRegistry.newCounter(
- GROUP, getMetricNameWithPrefix(TRANSFORM_IP_THROUGHPUT, transformName)));
- transformCacheSize.putIfAbsent(
- transformName,
- ctx.getTaskContext()
- .getTaskMetricsRegistry()
- .newGauge(
- GROUP,
- getMetricNameWithPrefix(TRANSFORM_ARRIVAL_TIME_CACHE_SIZE, transformName),
- 0L));
- }
-
- public Timer getTransformLatencyMetric(String transformName) {
- return transformLatency.get(transformName);
- }
-
- public Counter getTransformInputThroughput(String transformName) {
- return transformInputThroughput.get(transformName);
- }
-
- public Counter getTransformOutputThroughput(String transformName) {
- return transformOutputThroughPut.get(transformName);
- }
-
- public Gauge getTransformCacheSize(String transformName) {
- return transformCacheSize.get(transformName);
- }
-
- public Gauge getTransformWatermarkProgress(String transformName) {
- return transformWatermarkProgress.get(transformName);
- }
-
- // Customize in-memory window size for timer, default from samza is 5 mins which causes memory
- // pressure if a lot of timers are registered
- private static Timer getTimerWithCustomizedLookBackWindow(String transformName) {
- return new Timer(
- getMetricNameWithPrefix(TRANSFORM_LATENCY_METRIC, transformName),
- new SlidingTimeWindowReservoir(DEFAULT_LOOKBACK_TIMER_WINDOW_SIZE_MS));
- }
-
- private static String getMetricNameWithPrefix(String metricName, String transformName) {
- // Replace all non-alphanumeric characters with underscore
- final String samzaSafeMetricName = transformName.replaceAll("[^A-Za-z0-9_]", "_");
- return String.format("%s-%s", samzaSafeMetricName, metricName);
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/package-info.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/package-info.java
deleted file mode 100644
index 97415846e310..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/package-info.java
+++ /dev/null
@@ -1,20 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-/** Internal implementation of the Beam runner for Apache Samza. */
-package org.apache.beam.runners.samza.metrics;
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/package-info.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/package-info.java
deleted file mode 100644
index 549a4d81c2de..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/package-info.java
+++ /dev/null
@@ -1,20 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-/** Internal implementation of the Beam runner for Apache Samza. */
-package org.apache.beam.runners.samza;
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/AsyncDoFnRunner.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/AsyncDoFnRunner.java
deleted file mode 100644
index e1bc9251a304..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/AsyncDoFnRunner.java
+++ /dev/null
@@ -1,188 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.runtime;
-
-import java.util.Collection;
-import java.util.Collections;
-import java.util.Map;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ExecutorService;
-import java.util.stream.Collectors;
-import org.apache.beam.runners.core.DoFnRunner;
-import org.apache.beam.runners.samza.SamzaPipelineOptions;
-import org.apache.beam.sdk.state.TimeDomain;
-import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.values.CausedByDrain;
-import org.apache.beam.sdk.values.KV;
-import org.apache.beam.sdk.values.WindowedValue;
-import org.checkerframework.checker.nullness.qual.Nullable;
-import org.joda.time.Instant;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * This {@link DoFnRunner} adds the capability of executing the {@link
- * org.apache.beam.sdk.transforms.DoFn.ProcessElement} in the thread pool, and returns the future to
- * the collector for the underlying async execution.
- */
-public class AsyncDoFnRunner implements DoFnRunner {
- private static final Logger LOG = LoggerFactory.getLogger(AsyncDoFnRunner.class);
-
- // A dummy key to represent null keys
- private static final Object NULL_KEY = new Object();
-
- private final DoFnRunner underlying;
- private final ExecutorService executor;
- private final OpEmitter emitter;
- private final FutureCollector futureCollector;
- private final boolean isStateful;
-
- /**
- * This map keeps track of the last outputFutures for a certain key. When the next element of the
- * key comes in, its outputFutures will be chained from the last outputFutures in the map. When
- * all futures of a key have been complete, the key entry will be removed. The map is bounded by
- * (bundle size * 2).
- */
- private final Map>>> keyedOutputFutures;
-
- public static AsyncDoFnRunner create(
- DoFnRunner runner,
- OpEmitter emitter,
- FutureCollector futureCollector,
- boolean isStateful,
- SamzaPipelineOptions options) {
-
- LOG.info("Run DoFn with {}", AsyncDoFnRunner.class.getName());
- return new AsyncDoFnRunner<>(runner, emitter, futureCollector, isStateful, options);
- }
-
- private AsyncDoFnRunner(
- DoFnRunner runner,
- OpEmitter emitter,
- FutureCollector futureCollector,
- boolean isStateful,
- SamzaPipelineOptions options) {
- this.underlying = runner;
- this.executor = options.getExecutorServiceForProcessElement();
- this.emitter = emitter;
- this.futureCollector = futureCollector;
- this.isStateful = isStateful;
- this.keyedOutputFutures = new ConcurrentHashMap<>();
- }
-
- @Override
- public void startBundle() {
- underlying.startBundle();
- }
-
- @Override
- public void processElement(WindowedValue elem) {
- final CompletableFuture>> outputFutures =
- isStateful ? processStateful(elem) : processElement(elem, null);
-
- futureCollector.addAll(outputFutures);
- }
-
- private CompletableFuture>> processElement(
- WindowedValue elem,
- @Nullable CompletableFuture>> prevOutputFuture) {
-
- final CompletableFuture>> prevFuture =
- prevOutputFuture == null
- ? CompletableFuture.completedFuture(Collections.emptyList())
- : prevOutputFuture;
-
- // For ordering by key, we chain the processing of the elem to the completion of
- // the previous output of the same key
- return prevFuture.thenApplyAsync(
- x -> {
- underlying.processElement(elem);
-
- return emitter.collectOutput().stream()
- .map(OpMessage::getElement)
- .collect(Collectors.toList());
- },
- executor);
- }
-
- private CompletableFuture>> processStateful(
- WindowedValue elem) {
- final Object key = getKey(elem);
-
- final CompletableFuture>> outputFutures =
- processElement(elem, keyedOutputFutures.get(key));
-
- // Update the latest outputFuture for key
- keyedOutputFutures.put(key, outputFutures);
-
- // Remove the outputFuture from the map once it's complete.
- // This ensures the map will be cleaned up immediately.
- return outputFutures.thenApply(
- output -> {
- // Under the condition that the outputFutures has not been updated
- keyedOutputFutures.remove(key, outputFutures);
- return output;
- });
- }
-
- /** Package private for testing. */
- boolean hasOutputFuturesForKey(Object key) {
- return keyedOutputFutures.containsKey(key);
- }
-
- @Override
- public void onTimer(
- String timerId,
- String timerFamilyId,
- KeyT key,
- BoundedWindow window,
- Instant timestamp,
- Instant outputTimestamp,
- TimeDomain timeDomain,
- CausedByDrain causedByDrain) {
- underlying.onTimer(
- timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, causedByDrain);
- }
-
- @Override
- public void finishBundle() {
- underlying.finishBundle();
- }
-
- @Override
- public void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {
- underlying.onWindowExpiration(window, timestamp, key);
- }
-
- @Override
- public DoFn getFn() {
- return underlying.getFn();
- }
-
- private Object getKey(WindowedValue elem) {
- KV kv = (KV) elem.getValue();
- if (kv == null) {
- return NULL_KEY;
- } else {
- Object key = kv.getKey();
- return key == null ? NULL_KEY : key;
- }
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/BundleManager.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/BundleManager.java
deleted file mode 100644
index 36ba19d7da3c..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/BundleManager.java
+++ /dev/null
@@ -1,90 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.runtime;
-
-import org.joda.time.Instant;
-
-/**
- * Bundle management for the {@link DoFnOp} that handles lifecycle of a bundle. It also serves as a
- * proxy for the {@link DoFnOp} to process watermark and decides to 1. Hold watermark if there is at
- * least one bundle in progress. 2. Propagates the watermark to downstream DAG, if all the previous
- * bundles have completed.
- *
- * A bundle is considered complete only when the outputs corresponding to each element in the
- * bundle have been resolved and the watermark associated with the bundle(if any) is propagated
- * downstream. The output of an element is considered resolved based on the nature of the ParDoFn 1.
- * In case of synchronous ParDo, outputs of the element is resolved immediately after the
- * processElement returns. 2. In case of asynchronous ParDo, outputs of the element is resolved when
- * all the future emitted by the processElement is resolved.
- *
- * @param output type of the {@link DoFnOp}
- */
-public interface BundleManager {
- /** Starts a new bundle if not already started, then adds an element to the existing bundle. */
- void tryStartBundle();
-
- /**
- * Signals a watermark event arrived. The BundleManager will decide if the watermark needs to be
- * processed, and notify the listener if needed.
- *
- * @param watermark
- * @param emitter
- */
- void processWatermark(Instant watermark, OpEmitter emitter);
-
- /**
- * Signals the BundleManager that a timer is up.
- *
- * @param keyedTimerData
- * @param emitter
- */
- void processTimer(KeyedTimerData keyedTimerData, OpEmitter emitter);
-
- /**
- * Fails the current bundle, throws away the pending output, and resets the bundle to an empty
- * state.
- *
- * @param t the throwable that caused the failure.
- */
- void signalFailure(Throwable t);
-
- /**
- * Tries to close the bundle, and reset the bundle to an empty state.
- *
- * @param emitter
- */
- void tryFinishBundle(OpEmitter emitter);
-
- /**
- * A listener used to track the lifecycle of a bundle. Typically, the lifecycle of a bundle
- * consists of 1. Start bundle - Invoked when the bundle is started 2. Finish bundle - Invoked
- * when the bundle is complete. Refer to the docs under {@link BundleManager} for definition on
- * when a bundle is considered complete. 3. onWatermark - Invoked when watermark is ready to be
- * propagated to downstream DAG. Refer to the docs under {@link BundleManager} on when watermark
- * is held vs propagated.
- *
- * @param
- */
- interface BundleProgressListener {
- void onBundleStarted();
-
- void onBundleFinished(OpEmitter emitter);
-
- void onWatermark(Instant watermark, OpEmitter emitter);
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/ClassicBundleManager.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/ClassicBundleManager.java
deleted file mode 100644
index 53b6968e1119..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/ClassicBundleManager.java
+++ /dev/null
@@ -1,327 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.runtime;
-
-import java.util.Collection;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.CompletionStage;
-import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.concurrent.atomic.AtomicLong;
-import java.util.concurrent.atomic.AtomicReference;
-import java.util.function.BiConsumer;
-import javax.annotation.Nullable;
-import org.apache.beam.runners.core.StateNamespaces;
-import org.apache.beam.runners.core.TimerInternals;
-import org.apache.beam.sdk.state.TimeDomain;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.values.CausedByDrain;
-import org.apache.beam.sdk.values.WindowedValue;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
-import org.apache.samza.operators.Scheduler;
-import org.joda.time.Duration;
-import org.joda.time.Instant;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * {@inheritDoc} Implementation of BundleManager for non-portable mode. Keeps track of the async
- * function completions.
- *
- * This class is not thread safe and the current implementation relies on the assumption that
- * messages are dispatched to BundleManager in a single threaded mode.
- */
-@SuppressWarnings({
- "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
-public class ClassicBundleManager implements BundleManager {
- private static final Logger LOG = LoggerFactory.getLogger(ClassicBundleManager.class);
- private static final long MIN_BUNDLE_CHECK_TIME_MS = 10L;
-
- private final long maxBundleSize;
- private final long maxBundleTimeMs;
- private final BundleProgressListener bundleProgressListener;
- private final FutureCollector futureCollector;
- private final Scheduler> bundleTimerScheduler;
- private final String bundleCheckTimerId;
-
- // Number elements belonging to the current active bundle
- private transient AtomicLong currentBundleElementCount;
- // Number of bundles that are in progress but not yet finished
- private transient AtomicLong pendingBundleCount;
- // Denotes the start time of the current active bundle
- private transient AtomicLong bundleStartTime;
- // Denotes if there is an active in progress bundle. Note at a given time, we can have multiple
- // bundle in progress.
- // This flag denotes if there is a bundle that is current and hasn't been closed.
- private transient AtomicBoolean isBundleStarted;
- // Holder for watermark which gets propagated when the bundle is finished.
- private transient Instant bundleWatermarkHold;
- // A future that is completed once all futures belonging to the current active bundle are
- // completed. The value is null if there are no futures in the current active bundle.
- private transient AtomicReference> currentActiveBundleDoneFutureReference;
- private transient CompletionStage watermarkFuture;
-
- public ClassicBundleManager(
- BundleProgressListener bundleProgressListener,
- FutureCollector futureCollector,
- long maxBundleSize,
- long maxBundleTimeMs,
- Scheduler> bundleTimerScheduler,
- String bundleCheckTimerId) {
- this.maxBundleSize = maxBundleSize;
- this.maxBundleTimeMs = maxBundleTimeMs;
- this.bundleProgressListener = bundleProgressListener;
- this.bundleTimerScheduler = bundleTimerScheduler;
- this.bundleCheckTimerId = bundleCheckTimerId;
- this.futureCollector = futureCollector;
-
- if (maxBundleSize > 1) {
- scheduleNextBundleCheck();
- }
-
- // instance variable initialization for bundle tracking
- this.bundleStartTime = new AtomicLong(Long.MAX_VALUE);
- this.currentActiveBundleDoneFutureReference = new AtomicReference<>();
- this.currentBundleElementCount = new AtomicLong(0L);
- this.isBundleStarted = new AtomicBoolean(false);
- this.pendingBundleCount = new AtomicLong(0L);
- this.watermarkFuture = CompletableFuture.completedFuture(null);
- }
-
- /*
- * Schedule in processing time to check whether the current bundle should be closed. Note that
- * we only approximately achieve max bundle time by checking as frequent as half of the max bundle
- * time set by users. This would violate the max bundle time by up to half of it but should
- * acceptable in most cases (and cheaper than scheduling a timer at the beginning of every bundle).
- */
- private void scheduleNextBundleCheck() {
- final Instant nextBundleCheckTime =
- Instant.now().plus(Duration.millis(maxBundleTimeMs / 2 + MIN_BUNDLE_CHECK_TIME_MS));
- final TimerInternals.TimerData timerData =
- TimerInternals.TimerData.of(
- this.bundleCheckTimerId,
- StateNamespaces.global(),
- nextBundleCheckTime,
- nextBundleCheckTime,
- TimeDomain.PROCESSING_TIME,
- CausedByDrain.NORMAL);
- bundleTimerScheduler.schedule(
- new KeyedTimerData<>(new byte[0], null, timerData), nextBundleCheckTime.getMillis());
- }
-
- @Override
- public void tryStartBundle() {
- futureCollector.prepare();
-
- if (isBundleStarted.compareAndSet(false, true)) {
- LOG.debug("Starting a new bundle.");
- // make sure the previous bundle is sealed and futures are cleared
- Preconditions.checkArgument(
- currentActiveBundleDoneFutureReference.get() == null,
- "Current active bundle done future should be null before starting a new bundle.");
- bundleStartTime.set(System.currentTimeMillis());
- pendingBundleCount.incrementAndGet();
- bundleProgressListener.onBundleStarted();
- }
-
- currentBundleElementCount.incrementAndGet();
- }
-
- @Override
- public void processWatermark(Instant watermark, OpEmitter emitter) {
- // propagate watermark immediately if no bundle is in progress and all the previous bundles have
- // completed.
- if (!isBundleStarted() && pendingBundleCount.get() == 0) {
- LOG.debug("Propagating watermark: {} directly since no bundle in progress.", watermark);
- bundleProgressListener.onWatermark(watermark, emitter);
- return;
- }
-
- // hold back the watermark since there is either a bundle in progress or previously closed
- // bundles are unfinished.
- this.bundleWatermarkHold = watermark;
-
- // for batch mode, the max watermark should force the bundle to close
- if (BoundedWindow.TIMESTAMP_MAX_VALUE.equals(watermark)) {
- /*
- * Due to lack of async watermark function, we block on the previous watermark futures before propagating the watermark
- * downstream. If a bundle is in progress tryFinishBundle() fill force the bundle to close and emit watermark.
- * If no bundle in progress, we progress watermark explicitly after the completion of previous watermark futures.
- */
- if (isBundleStarted()) {
- LOG.info(
- "Received max watermark. Triggering finish bundle before flushing the watermark downstream.");
- tryFinishBundle(emitter);
- watermarkFuture.toCompletableFuture().join();
- } else {
- LOG.info(
- "Received max watermark. Waiting for previous bundles to complete before flushing the watermark downstream.");
- watermarkFuture.toCompletableFuture().join();
- bundleProgressListener.onWatermark(watermark, emitter);
- }
- }
- }
-
- @Override
- public void processTimer(KeyedTimerData keyedTimerData, OpEmitter emitter) {
- // this is internal timer in processing time to check whether a bundle should be closed
- if (bundleCheckTimerId.equals(keyedTimerData.getTimerData().getTimerId())) {
- tryFinishBundle(emitter);
- scheduleNextBundleCheck();
- }
- }
-
- /**
- * Signal the bundle manager to handle failure. We discard the output collected as part of
- * processing the current element and reset the bundle count.
- *
- * @param t failure cause
- */
- @Override
- public void signalFailure(Throwable t) {
- LOG.error("Encountered error during processing the message. Discarding the output due to: ", t);
- futureCollector.discard();
- // reset the bundle start flag only if the bundle has started
- isBundleStarted.compareAndSet(true, false);
-
- // bundle start may not necessarily mean we have actually started the bundle since some of the
- // invariant check conditions within bundle start could throw exceptions. so rely on bundle
- // start time
- if (bundleStartTime.get() != Long.MAX_VALUE) {
- currentBundleElementCount.set(0L);
- bundleStartTime.set(Long.MAX_VALUE);
- pendingBundleCount.decrementAndGet();
- currentActiveBundleDoneFutureReference.set(null);
- }
- }
-
- @Override
- public void tryFinishBundle(OpEmitter emitter) {
-
- // we need to seal the output for each element within a bundle irrespective of the whether we
- // decide to finish the
- // bundle or not
- CompletionStage>> outputFuture = futureCollector.finish();
-
- if (shouldFinishBundle() && isBundleStarted.compareAndSet(true, false)) {
- LOG.debug("Finishing the current bundle.");
-
- // reset the bundle count
- // seal the bundle and emit the result future (collection of results)
- // chain the finish bundle invocation on the finish bundle
- currentBundleElementCount.set(0L);
- bundleStartTime.set(Long.MAX_VALUE);
- Instant watermarkHold = bundleWatermarkHold;
- bundleWatermarkHold = null;
-
- CompletionStage currentActiveBundleDoneFuture =
- currentActiveBundleDoneFutureReference.get();
- outputFuture =
- outputFuture.thenCombine(
- currentActiveBundleDoneFuture != null
- ? currentActiveBundleDoneFuture
- : CompletableFuture.completedFuture(null),
- (res, ignored) -> {
- bundleProgressListener.onBundleFinished(emitter);
- return res;
- });
-
- BiConsumer>, Void> watermarkPropagationFn;
- if (watermarkHold == null) {
- watermarkPropagationFn = (ignored, res) -> pendingBundleCount.decrementAndGet();
- } else {
- watermarkPropagationFn =
- (ignored, res) -> {
- LOG.debug("Propagating watermark: {} to downstream.", watermarkHold);
- bundleProgressListener.onWatermark(watermarkHold, emitter);
- pendingBundleCount.decrementAndGet();
- };
- }
-
- // We chain the current watermark emission with previous watermark and the output futures
- // since bundles can finish out of order but we still want the watermark to be emitted in
- // order.
- watermarkFuture = outputFuture.thenAcceptBoth(watermarkFuture, watermarkPropagationFn);
- currentActiveBundleDoneFutureReference.set(null);
- } else if (isBundleStarted.get()) {
- final CompletableFuture>> finalOutputFuture =
- outputFuture.toCompletableFuture();
- currentActiveBundleDoneFutureReference.updateAndGet(
- maybePrevFuture -> {
- CompletableFuture prevFuture =
- maybePrevFuture != null ? maybePrevFuture : CompletableFuture.completedFuture(null);
-
- return CompletableFuture.allOf(prevFuture, finalOutputFuture);
- });
- }
-
- // emit the future to the propagate it to rest of the DAG
- emitter.emitFuture(outputFuture);
- }
-
- @VisibleForTesting
- long getCurrentBundleElementCount() {
- return currentBundleElementCount.longValue();
- }
-
- @VisibleForTesting
- @Nullable
- CompletionStage getCurrentBundleDoneFuture() {
- return currentActiveBundleDoneFutureReference.get();
- }
-
- @VisibleForTesting
- void setCurrentBundleDoneFuture(CompletableFuture currentBundleResultFuture) {
- this.currentActiveBundleDoneFutureReference.set(currentBundleResultFuture);
- }
-
- @VisibleForTesting
- long getPendingBundleCount() {
- return pendingBundleCount.longValue();
- }
-
- @VisibleForTesting
- void setPendingBundleCount(long value) {
- pendingBundleCount.set(value);
- }
-
- @VisibleForTesting
- boolean isBundleStarted() {
- return isBundleStarted.get();
- }
-
- @VisibleForTesting
- void setBundleWatermarkHold(Instant watermark) {
- this.bundleWatermarkHold = watermark;
- }
-
- /**
- * We close the current bundle in progress if one of the following criteria is met 1. The bundle
- * count ≥ maxBundleSize 2. Time elapsed since the bundle started is ≥ maxBundleTimeMs 3.
- * Watermark hold equals to TIMESTAMP_MAX_VALUE which usually is the case for bounded jobs
- *
- * @return true - if one of the criteria above is satisfied; false - otherwise
- */
- private boolean shouldFinishBundle() {
- return isBundleStarted.get()
- && (currentBundleElementCount.get() >= maxBundleSize
- || System.currentTimeMillis() - bundleStartTime.get() >= maxBundleTimeMs
- || BoundedWindow.TIMESTAMP_MAX_VALUE.equals(bundleWatermarkHold));
- }
-}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java
deleted file mode 100644
index bc87e2460ec4..000000000000
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java
+++ /dev/null
@@ -1,579 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.samza.runtime;
-
-import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
-
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.ServiceLoader;
-import java.util.concurrent.CompletionStage;
-import java.util.function.Function;
-import org.apache.beam.model.pipeline.v1.RunnerApi;
-import org.apache.beam.runners.core.DoFnRunner;
-import org.apache.beam.runners.core.PushbackSideInputDoFnRunner;
-import org.apache.beam.runners.core.SideInputHandler;
-import org.apache.beam.runners.core.SimplePushbackSideInputDoFnRunner;
-import org.apache.beam.runners.core.StateNamespace;
-import org.apache.beam.runners.core.StateNamespaces;
-import org.apache.beam.runners.core.TimerInternals;
-import org.apache.beam.runners.fnexecution.control.ExecutableStageContext;
-import org.apache.beam.runners.fnexecution.control.StageBundleFactory;
-import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
-import org.apache.beam.runners.samza.SamzaExecutionContext;
-import org.apache.beam.runners.samza.SamzaPipelineOptions;
-import org.apache.beam.runners.samza.util.DoFnUtils;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.DoFnSchemaInformation;
-import org.apache.beam.sdk.transforms.join.RawUnionValue;
-import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
-import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
-import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
-import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.util.WindowedValueMultiReceiver;
-import org.apache.beam.sdk.util.construction.graph.ExecutableStage;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.sdk.values.WindowedValue;
-import org.apache.beam.sdk.values.WindowedValues;
-import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterators;
-import org.apache.samza.config.Config;
-import org.apache.samza.context.Context;
-import org.apache.samza.operators.Scheduler;
-import org.joda.time.Instant;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/** Samza operator for {@link DoFn}. */
-@SuppressWarnings({
- "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
- "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
-public class DoFnOp implements Op {
- private static final Logger LOG = LoggerFactory.getLogger(DoFnOp.class);
-
- private final TupleTag