diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/BoundedReadEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/BoundedReadEvaluatorFactory.java index b31798e7e46b..aeedc84eb3fe 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/BoundedReadEvaluatorFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/BoundedReadEvaluatorFactory.java @@ -17,6 +17,8 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import com.google.auto.value.AutoValue; import java.io.IOException; import java.util.Collection; @@ -52,8 +54,7 @@ * for the {@link PrimitiveBoundedRead SplittableParDo.PrimitiveBoundedRead} {@link PTransform}. */ @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) final class BoundedReadEvaluatorFactory implements TransformEvaluatorFactory { /** @@ -149,7 +150,8 @@ public void processElement(WindowedValue> element) BoundedSource source = element.getValue().getSource(); try (final BoundedReader reader = source.createReader(options)) { boolean contentsRemaining = reader.start(); - Future> residualFuture = startDynamicSplitThread(source, reader); + Future<@Nullable BoundedSource> residualFuture = + startDynamicSplitThread(source, reader); UncommittedBundle output = evaluationContext.createBundle(outputPCollection); while (contentsRemaining) { output.add( @@ -159,24 +161,24 @@ public void processElement(WindowedValue> element) } resultBuilder.addOutput(output); try { - BoundedSource residual = residualFuture.get(); + @Nullable BoundedSource residual = residualFuture.get(); if (residual != null) { resultBuilder.addUnprocessedElements( element.withValue(BoundedSourceShard.of(residual))); } } catch (ExecutionException exex) { // Un-and-rewrap the exception thrown by attempting to split - throw UserCodeException.wrap(exex.getCause()); + throw UserCodeException.wrap(checkStateNotNull(exex.getCause())); } } } - private Future> startDynamicSplitThread( + private Future<@Nullable BoundedSource> startDynamicSplitThread( BoundedSource source, BoundedReader reader) throws Exception { if (source.getEstimatedSizeBytes(options) > minimumDynamicSplitSize) { return produceSplitExecutor.submit(new GenerateSplitAtHalfwayPoint<>(reader)); } else { - SettableFuture> emptyFuture = SettableFuture.create(); + SettableFuture<@Nullable BoundedSource> emptyFuture = SettableFuture.create(); emptyFuture.set(null); return emptyFuture; } @@ -230,7 +232,8 @@ public Collection>> getInitialInputs( } } - private static class GenerateSplitAtHalfwayPoint implements Callable> { + private static class GenerateSplitAtHalfwayPoint + implements Callable<@Nullable BoundedSource> { private final BoundedReader reader; private GenerateSplitAtHalfwayPoint(BoundedReader reader) { @@ -238,7 +241,7 @@ private GenerateSplitAtHalfwayPoint(BoundedReader reader) { } @Override - public BoundedSource call() throws Exception { + public @Nullable BoundedSource call() throws Exception { // Splits at halfway of the remaining work. Double currentlyConsumed = reader.getFractionConsumed(); if (currentlyConsumed == null || currentlyConsumed == 1.0) { diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CloningBundleFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CloningBundleFactory.java index c3d30bfbf744..182200b021c0 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CloningBundleFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CloningBundleFactory.java @@ -17,6 +17,8 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import org.apache.beam.runners.local.StructuralKey; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderException; @@ -25,15 +27,13 @@ import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** * A {@link BundleFactory} where a created {@link UncommittedBundle} clones all elements added to it * using the coder of the {@link PCollection}. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) class CloningBundleFactory implements BundleFactory { private static final CloningBundleFactory INSTANCE = new CloningBundleFactory(); @@ -70,11 +70,12 @@ private static class CloningBundle implements UncommittedBundle { private CloningBundle(UncommittedBundle underlying) { this.underlying = underlying; - this.coder = underlying.getPCollection().getCoder(); + // CloningBundle only wraps non-root bundles, whose PCollection is always present. + this.coder = checkStateNotNull(underlying.getPCollection()).getCoder(); } @Override - public PCollection getPCollection() { + public @Nullable PCollection getPCollection() { return underlying.getPCollection(); } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CommittedResult.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CommittedResult.java index af0c92d09c51..0a58fdf9e717 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CommittedResult.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CommittedResult.java @@ -24,9 +24,6 @@ /** A {@link TransformResult} that has been committed. */ @AutoValue -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) abstract class CommittedResult { /** Returns the {@link AppliedPTransform} that produced this result. */ public abstract ExecutableT getExecutable(); diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CopyOnAccessInMemoryStateInternals.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CopyOnAccessInMemoryStateInternals.java index 720ce77c388d..2d11cdb74049 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CopyOnAccessInMemoryStateInternals.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CopyOnAccessInMemoryStateInternals.java @@ -56,6 +56,7 @@ import org.apache.beam.sdk.transforms.windowing.TimestampCombiner; import org.apache.beam.sdk.util.CombineFnUtil; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; @@ -65,10 +66,9 @@ * accessed, an independent copy will be created within this table. */ @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) -class CopyOnAccessInMemoryStateInternals implements StateInternals { +class CopyOnAccessInMemoryStateInternals implements StateInternals { private final CopyOnAccessInMemoryStateTable table; private K key; @@ -77,12 +77,13 @@ class CopyOnAccessInMemoryStateInternals implements StateInternals { * Creates a new {@link CopyOnAccessInMemoryStateInternals} with the underlying (possibly null) * StateInternals. */ - public static CopyOnAccessInMemoryStateInternals withUnderlying( + public static CopyOnAccessInMemoryStateInternals withUnderlying( K key, @Nullable CopyOnAccessInMemoryStateInternals underlying) { return new CopyOnAccessInMemoryStateInternals<>(key, underlying); } - private CopyOnAccessInMemoryStateInternals(K key, CopyOnAccessInMemoryStateInternals underlying) { + private CopyOnAccessInMemoryStateInternals( + K key, @Nullable CopyOnAccessInMemoryStateInternals underlying) { this.key = key; table = new CopyOnAccessInMemoryStateTable(underlying == null ? null : underlying.table); } @@ -169,7 +170,7 @@ private static class CopyOnAccessInMemoryStateTable extends StateTable { /** The earliest watermark hold in this table. */ private Optional earliestWatermarkHold; - public CopyOnAccessInMemoryStateTable(StateTable underlying) { + public CopyOnAccessInMemoryStateTable(@Nullable StateTable underlying) { this.underlying = Optional.ofNullable(underlying); binderFactory = new CopyOnBindBinderFactory(this.underlying); earliestWatermarkHold = Optional.empty(); diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectExecutionContext.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectExecutionContext.java index 249934c43084..63c86ef79df9 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectExecutionContext.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectExecutionContext.java @@ -28,6 +28,7 @@ import org.apache.beam.runners.direct.WatermarkManager.TransformWatermarks; import org.apache.beam.runners.local.StructuralKey; import org.apache.beam.sdk.transforms.DoFn.BundleFinalizer; +import org.checkerframework.checker.nullness.qual.Nullable; /** * Execution Context for the {@link DirectRunner}. @@ -36,21 +37,20 @@ * for each thread that requires it. */ @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) class DirectExecutionContext { private final Clock clock; private final StructuralKey key; - private final CopyOnAccessInMemoryStateInternals existingState; + private final @Nullable CopyOnAccessInMemoryStateInternals existingState; private final TransformWatermarks watermarks; private Map cachedStepContexts = new LinkedHashMap<>(); public DirectExecutionContext( Clock clock, StructuralKey key, - CopyOnAccessInMemoryStateInternals existingState, + @Nullable CopyOnAccessInMemoryStateInternals existingState, TransformWatermarks watermarks) { this.clock = clock; this.key = key; @@ -69,9 +69,9 @@ public DirectStepContext getStepContext(String stepName) { /** Step Context for the {@link DirectRunner}. */ public class DirectStepContext implements StepContext { - private CopyOnAccessInMemoryStateInternals stateInternals; - private DirectTimerInternals timerInternals; - private InMemoryBundleFinalizer bundleFinalizer; + private @Nullable CopyOnAccessInMemoryStateInternals stateInternals; + private @Nullable DirectTimerInternals timerInternals; + private @Nullable InMemoryBundleFinalizer bundleFinalizer; public DirectStepContext() {} @@ -111,7 +111,7 @@ public List getAndClearFinalizations() { * Commits the state of this step, and returns the committed state. If the step has not accessed * any state, return null. */ - public CopyOnAccessInMemoryStateInternals commitState() { + public @Nullable CopyOnAccessInMemoryStateInternals commitState() { if (stateInternals != null) { return stateInternals.commit(); } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraph.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraph.java index 00f246089493..3c3782e724ac 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraph.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraph.java @@ -17,6 +17,8 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import java.util.Collection; import java.util.List; import java.util.Map; @@ -34,7 +36,7 @@ * Methods for interacting with the underlying structure of a {@link Pipeline} that is being * executed with the {@link DirectRunner}. */ -@SuppressWarnings({"keyfor", "nullness"}) // TODO(https://github.com/apache/beam/issues/20497) +@SuppressWarnings({"keyfor"}) // TODO(https://github.com/apache/beam/issues/20497) class DirectGraph implements ExecutableGraph, PValue> { private final Map, AppliedPTransform> producers; private final Map, AppliedPTransform> viewWriters; @@ -68,7 +70,7 @@ private DirectGraph( @Override public AppliedPTransform getProducer(PValue produced) { if (produced instanceof PCollection) { - return producers.get(produced); + return checkStateNotNull(producers.get(produced), "No producer for %s", produced); } else if (produced instanceof PCollectionView) { return getWriter((PCollectionView) produced); } @@ -96,7 +98,7 @@ public Collection getPerElementInputs(AppliedPTransform transfo } private AppliedPTransform getWriter(PCollectionView view) { - return viewWriters.get(view); + return checkStateNotNull(viewWriters.get(view), "No writer for %s", view); } @Override @@ -119,6 +121,6 @@ Set> getViews() { } String getStepName(AppliedPTransform step) { - return stepNames.get(step); + return checkStateNotNull(stepNames.get(step), "No step name for %s", step); } } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraphVisitor.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraphVisitor.java index adaae0b5494d..436e013bbf23 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraphVisitor.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraphVisitor.java @@ -47,9 +47,6 @@ * {@link Pipeline}. This is used to schedule consuming {@link PTransform PTransforms} to consume * input after the upstream transform has produced and committed output. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) class DirectGraphVisitor extends PipelineVisitor.Defaults { private static final Logger LOG = LoggerFactory.getLogger(DirectGraphVisitor.class); @@ -119,12 +116,11 @@ public void visitPrimitiveTransform(TransformHierarchy.Node node) { allConsumers.put(value, appliedTransform); } } - if (node.getTransform() instanceof ParDo.MultiOutput) { - consumedViews.addAll( - ((ParDo.MultiOutput) node.getTransform()).getSideInputs().values()); - } else if (node.getTransform() instanceof WriteView) { - viewWriters.put( - ((WriteView) node.getTransform()).getView(), node.toAppliedPTransform(getPipeline())); + PTransform transform = node.getTransform(); + if (transform instanceof ParDo.MultiOutput) { + consumedViews.addAll(((ParDo.MultiOutput) transform).getSideInputs().values()); + } else if (transform instanceof WriteView) { + viewWriters.put(((WriteView) transform).getView(), node.toAppliedPTransform(getPipeline())); } } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectMetrics.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectMetrics.java index b7cf0357447a..dbbf6d925ff4 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectMetrics.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectMetrics.java @@ -47,12 +47,10 @@ import org.apache.beam.sdk.metrics.MetricsFilter; import org.apache.beam.sdk.metrics.StringSetResult; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; /** Implementation of {@link MetricResults} for the Direct Runner. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) class DirectMetrics extends MetricResults { private interface MetricAggregation { @@ -69,7 +67,7 @@ private interface MetricAggregation { * @param The type of raw data received and aggregated across updates. * @param The type of result extracted from the data. */ - private static class DirectMetric { + private static class DirectMetric { private final MetricAggregation aggregation; private final Executor executor; @@ -326,10 +324,10 @@ public MetricQueryResults queryMetrics(@Nullable MetricsFilter filter) { } private void maybeExtractResult( - MetricsFilter filter, + @Nullable MetricsFilter filter, ImmutableList.Builder> resultsBuilder, Map.Entry> entry) { - if (MetricFiltering.matches(filter, entry.getKey())) { + if (filter == null || MetricFiltering.matches(filter, entry.getKey())) { resultsBuilder.add( MetricResult.create( entry.getKey(), diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java index 6997b590e017..9cc94ceaea9f 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java @@ -17,6 +17,8 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; import java.util.Collection; @@ -66,9 +68,6 @@ * contained within a {@link Pipeline} does not break assumptions within the Beam model, to improve * the ability to execute a {@link Pipeline} at scale on a distributed backend. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) public class DirectRunner extends PipelineRunner { enum Enforcement { @@ -222,7 +221,7 @@ public DirectPipelineResult run(Pipeline pipeline) { try { result.waitUntilFinish(); } catch (UserCodeException userException) { - throw new PipelineExecutionException(userException.getCause()); + throw new PipelineExecutionException(checkStateNotNull(userException.getCause())); } catch (RuntimeException | OutOfMemoryError e) { throw e; } catch (Throwable t) { @@ -368,6 +367,10 @@ public State cancel() { * org.apache.beam.sdk.PipelineResult.State#FAILED}. */ @Override + // PipelineResult.waitUntilFinish(Duration) is documented to return null on timeout, but its + // interface return type is not annotated @Nullable; annotating it would ripple across every + // runner's override, so the timeout null is suppressed here rather than in the sdk-core API. + @SuppressWarnings("nullness") public State waitUntilFinish(Duration duration) { if (this.state.isTerminal()) { return this.state; @@ -379,7 +382,7 @@ public State waitUntilFinish(Duration duration) { // Emulates the behavior of Pipeline#run(), where a stack trace caused by a // UserCodeException is truncated and replaced with the stack starting at the call to // waitToFinish - throw new Pipeline.PipelineExecutionException(uce.getCause()); + throw new Pipeline.PipelineExecutionException(checkStateNotNull(uce.getCause())); } catch (Exception e) { if (e instanceof InterruptedException) { Thread.currentThread().interrupt(); diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectWriteViewVisitor.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectWriteViewVisitor.java index 8bb168097853..e7edeaae7ca7 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectWriteViewVisitor.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectWriteViewVisitor.java @@ -17,6 +17,8 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import java.util.HashSet; import java.util.Set; import org.apache.beam.sdk.Pipeline; @@ -41,15 +43,12 @@ * Adds a {@link DirectRunner}-specific {@link WriteView} step for each {@link PCollectionView} for * scheduling materialization of side inputs. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) class DirectWriteViewVisitor extends PipelineVisitor.Defaults { /** Private URN for identifying {@link DirectRunner}-specific view writing transform. */ static final String DIRECT_WRITE_VIEW_URN = "beam:directrunner:transforms:write_view:v1"; - private Set> viewsToWrite; + private Set> viewsToWrite = new HashSet<>(); @Override public void enterPipeline(Pipeline p) { @@ -66,14 +65,16 @@ public void leavePipeline(Pipeline p) { @Override public void visitPrimitiveTransform(TransformHierarchy.Node node) { - if (node.getTransform() instanceof ParDo.MultiOutput) { - ParDo.MultiOutput parDo = (ParDo.MultiOutput) node.getTransform(); + PTransform transform = node.getTransform(); + if (transform instanceof ParDo.MultiOutput) { + ParDo.MultiOutput parDo = (ParDo.MultiOutput) transform; viewsToWrite.addAll(parDo.getSideInputs().values()); } } private void visitView(PCollectionView view) { - PCollection collectionToMaterialize = (PCollection) view.getPCollection(); + PCollection collectionToMaterialize = + (PCollection) checkStateNotNull(view.getPCollection()); collectionToMaterialize.apply("GroupAndWriteView", new GroupAndWriteView<>(view)); } @@ -95,7 +96,8 @@ public PCollection expand(final PCollection input) { iterable = input .apply( - MapElements.into(TypeDescriptors.iterables(input.getTypeDescriptor())) + MapElements.into( + TypeDescriptors.iterables(checkStateNotNull(input.getTypeDescriptor()))) .via(Lists::newArrayList)) .setCoder(IterableCoder.of(input.getCoder())); } else { diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DisplayDataValidator.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DisplayDataValidator.java index 69bb98bf8ba0..b7a94dc666a2 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DisplayDataValidator.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DisplayDataValidator.java @@ -22,14 +22,12 @@ import org.apache.beam.sdk.runners.TransformHierarchy; import org.apache.beam.sdk.transforms.display.DisplayData; import org.apache.beam.sdk.transforms.display.HasDisplayData; +import org.checkerframework.checker.nullness.qual.Nullable; /** * Validate correct implementation of {@link DisplayData} by evaluating {@link * HasDisplayData#populateDisplayData(DisplayData.Builder)} during pipeline construction. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) class DisplayDataValidator { // Do not instantiate private DisplayDataValidator() {} @@ -46,8 +44,11 @@ private static void validateTransforms(Pipeline pipeline) { pipeline.traverseTopologically(Visitor.INSTANCE); } - private static void evaluateDisplayData(HasDisplayData component) { - DisplayData.from(component); + private static void evaluateDisplayData(@Nullable HasDisplayData component) { + // A node's transform is null only for the root node, which has no display data to validate. + if (component != null) { + DisplayData.from(component); + } } private static class Visitor extends Pipeline.PipelineVisitor.Defaults { diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/EvaluationContext.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/EvaluationContext.java index 0fca686bced7..5920225d50f6 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/EvaluationContext.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/EvaluationContext.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; import java.util.Collection; @@ -51,6 +52,7 @@ 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.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.MoreExecutors; +import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.joda.time.Instant; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -73,9 +75,8 @@ */ @SuppressWarnings({ "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "keyfor", - "nullness" -}) // TODO(https://github.com/apache/beam/issues/20497) + "keyfor" +}) class EvaluationContext { private static final Logger LOG = LoggerFactory.getLogger(EvaluationContext.class); @@ -97,7 +98,9 @@ class EvaluationContext { private final ConcurrentMap applicationStateInternals; - private final SideInputContainer sideInputContainer; + // Set once in create() after construction completes, because SideInputContainer.create needs a + // fully-initialized EvaluationContext (it holds the context to schedule callbacks later). + private @MonotonicNonNull SideInputContainer sideInputContainer; private final DirectMetrics metrics; @@ -109,7 +112,11 @@ public static EvaluationContext create( DirectGraph graph, Set keyedPValues, ExecutorService executorService) { - return new EvaluationContext(clock, bundleFactory, graph, keyedPValues, executorService); + EvaluationContext evaluationContext = + new EvaluationContext(clock, bundleFactory, graph, keyedPValues, executorService); + evaluationContext.sideInputContainer = + SideInputContainer.create(evaluationContext, graph.getViews()); + return evaluationContext; } private EvaluationContext( @@ -124,7 +131,6 @@ private EvaluationContext( this.keyedPValues = keyedPValues; this.watermarkManager = WatermarkManager.create(clock, graph, AppliedPTransform::getFullName); - this.sideInputContainer = SideInputContainer.create(this, graph.getViews()); this.applicationStateInternals = new ConcurrentHashMap<>(); this.metrics = new DirectMetrics(executorService); @@ -222,7 +228,9 @@ private Iterable> commitBundles( Iterable> bundles) { ImmutableList.Builder> completed = ImmutableList.builder(); for (UncommittedBundle inProgress : bundles) { - AppliedPTransform producing = graph.getProducer(inProgress.getPCollection()); + // Bundles committed here are transform outputs, which always target a PCollection. + AppliedPTransform producing = + graph.getProducer(checkStateNotNull(inProgress.getPCollection())); TransformWatermarks watermarks = watermarkManager.getWatermarks(producing); CommittedBundle committed = inProgress.commit(watermarks.getSynchronizedProcessingOutputTime()); @@ -283,7 +291,7 @@ public boolean isKeyed(PValue pValue) { */ public PCollectionViewWriter createPCollectionViewWriter( PCollection> input, final PCollectionView output) { - return values -> sideInputContainer.write(output, values); + return values -> checkStateNotNull(sideInputContainer).write(output, values); } /** @@ -371,7 +379,7 @@ String getStepName(AppliedPTransform application) { */ public ReadyCheckingSideInputReader createSideInputReader( final List> sideInputs) { - return sideInputContainer.createReaderForViews(sideInputs); + return checkStateNotNull(sideInputContainer).createReaderForViews(sideInputs); } /** Returns the metrics container for this pipeline. */ diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ExecutorServiceParallelExecutor.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ExecutorServiceParallelExecutor.java index 603832843bb6..dd4df5df353d 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ExecutorServiceParallelExecutor.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ExecutorServiceParallelExecutor.java @@ -60,9 +60,6 @@ * An {@link PipelineExecutor} that uses an underlying {@link ExecutorService} and {@link * EvaluationContext} to execute a {@link Pipeline}. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) final class ExecutorServiceParallelExecutor implements PipelineExecutor, BundleProcessor, CommittedBundle, AppliedPTransform> { @@ -122,7 +119,7 @@ private ExecutorServiceParallelExecutor( CacheBuilder.newBuilder() .weakValues() .removalListener(shutdownExecutorServiceListener()) - .build(serialTransformExecutorServiceCacheLoader()); + .build(serialTransformExecutorServiceCacheLoader(executorService)); this.visibleUpdates = new QueueMessageReceiver(); @@ -130,8 +127,8 @@ private ExecutorServiceParallelExecutor( executorFactory = new DirectTransformExecutor.Factory(context, registry, transformEnforcements); } - private CacheLoader - serialTransformExecutorServiceCacheLoader() { + private static CacheLoader + serialTransformExecutorServiceCacheLoader(ExecutorService executorService) { return new CacheLoader() { @Override public TransformExecutorService load(StepAndKey stepAndKey) throws Exception { @@ -140,7 +137,8 @@ public TransformExecutorService load(StepAndKey stepAndKey) throws Exception { }; } - private RemovalListener shutdownExecutorServiceListener() { + private static RemovalListener + shutdownExecutorServiceListener() { return notification -> { TransformExecutorService service = notification.getValue(); if (service != null) { @@ -168,10 +166,11 @@ public void start(DirectGraph graph, RootProviderRegistry rootProviderRegistry) throw UserCodeException.wrap(e); } finally { // Metrics emitted initial split are reported along with the first bundle - if (pending.peek() != null) { + CommittedBundle firstPending = pending.peek(); + if (firstPending != null) { evaluationContext .getMetrics() - .commitPhysical(pending.peek(), metricsContainer.getCumulative()); + .commitPhysical(firstPending, metricsContainer.getCumulative()); } } pendingRootBundles.put(root, pending); @@ -224,7 +223,8 @@ private void evaluateBundle( final CompletionCallback onComplete) { TransformExecutorService transformExecutor; - if (isKeyed(bundle.getPCollection())) { + @Nullable PCollection pcollection = bundle.getPCollection(); + if (pcollection != null && isKeyed(pcollection)) { final StepAndKey stepAndKey = StepAndKey.of(transform, bundle.getKey()); // This executor will remain reachable until it has executed all scheduled transforms. // The TransformExecutors keep a strong reference to the Executor, the ExecutorService keeps @@ -248,7 +248,7 @@ private boolean isKeyed(PValue pvalue) { } @Override - public State waitUntilFinish(Duration duration) throws Exception { + public @Nullable State waitUntilFinish(Duration duration) throws Exception { Instant completionTime; if (duration.equals(Duration.ZERO)) { completionTime = new Instant(Long.MAX_VALUE); @@ -259,7 +259,7 @@ public State waitUntilFinish(Duration duration) throws Exception { while (Instant.now().isBefore(completionTime)) { // Get an update; don't block forever if another thread has handled it. The call to poll will // wait the entire timeout; this call primarily exists to relinquish any core. - VisibleExecutorUpdate update = visibleUpdates.tryNext(Duration.millis(25L)); + @Nullable VisibleExecutorUpdate update = visibleUpdates.tryNext(Duration.millis(25L)); if (update == null && pipelineState.get().isTerminal()) { // state and updates have seperate locks so it is possible for an update @@ -300,7 +300,8 @@ public State getPipelineState() { } private boolean isTerminalStateUpdate(VisibleExecutorUpdate update) { - return update.getNewState() != null && update.getNewState().isTerminal(); + State newState = update.getNewState(); + return newState != null && newState.isTerminal(); } @Override @@ -348,7 +349,7 @@ private void shutdownIfNecessary(State newState) { } catch (final Exception e) { errors.add(e); } - IllegalStateException exception = null; + @Nullable IllegalStateException exception = null; try { if (!errors.isEmpty()) { exception = @@ -393,12 +394,12 @@ public static VisibleExecutorUpdate cancelled() { return new VisibleExecutorUpdate(State.CANCELLED, null); } - private VisibleExecutorUpdate(State newState, @Nullable Throwable exception) { + private VisibleExecutorUpdate(@Nullable State newState, @Nullable Throwable exception) { this.thrown = Optional.ofNullable(exception); this.newState = newState; } - State getNewState() { + @Nullable State getNewState() { return newState; } } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java index c6726fb3463f..003fc3d090dc 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java @@ -17,6 +17,8 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import java.util.ArrayList; import java.util.Collection; import java.util.stream.Collectors; @@ -57,9 +59,8 @@ */ @SuppressWarnings({ "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "keyfor", - "nullness" -}) // TODO(https://github.com/apache/beam/issues/20497) + "keyfor" +}) class GroupAlsoByWindowEvaluatorFactory implements TransformEvaluatorFactory { private final EvaluationContext evaluationContext; private final PipelineOptions options; @@ -145,7 +146,9 @@ public GroupAlsoByWindowEvaluator( unprocessedElements = ImmutableList.builder(); Coder valueCoder = - application.getTransform().getValueCoder(inputBundle.getPCollection().getCoder()); + application + .getTransform() + .getValueCoder(checkStateNotNull(inputBundle.getPCollection()).getCoder()); reduceFn = SystemReduceFn.buffering(valueCoder); droppedDueToLateness = Metrics.counter( @@ -190,8 +193,8 @@ public void processElement(WindowedValue> element) throws Ex @Override public TransformResult> finishBundle() throws Exception { - // State is initialized within the constructor. It can never be null. - CopyOnAccessInMemoryStateInternals state = stepContext.commitState(); + // processElement always accesses state internals, so commitState is non-null here. + CopyOnAccessInMemoryStateInternals state = checkStateNotNull(stepContext.commitState()); return StepTransformResult.>withHold( application, state.getEarliestWatermarkHold()) .withState(state) @@ -226,7 +229,7 @@ Iterable> dropExpiredWindows( + "window: {} since it is too far behind inputWatermark: {}", DirectGroupAlsoByWindow.class.getSimpleName(), input.getTimestamp(), - key, + String.valueOf(key), window, timerInternals.currentInputWatermarkTime()); } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupByKeyOnlyEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupByKeyOnlyEvaluatorFactory.java index 121c6f8bd5d9..e73828904cfa 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupByKeyOnlyEvaluatorFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupByKeyOnlyEvaluatorFactory.java @@ -44,8 +44,7 @@ * PTransform}. */ @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) class GroupByKeyOnlyEvaluatorFactory implements TransformEvaluatorFactory { private final EvaluationContext evaluationContext; @@ -101,13 +100,13 @@ public GroupByKeyOnlyEvaluator( this.groupingMap = new HashMap<>(); } - private Coder getKeyCoder(Coder> coder) { + private static Coder getKeyCoder(Coder> coder) { checkState( coder instanceof KvCoder, "%s requires a coder of class %s." + " This is an internal error; this is checked during pipeline construction" + " but became corrupted.", - getClass().getSimpleName(), + GroupByKeyOnlyEvaluator.class.getSimpleName(), KvCoder.class.getSimpleName()); @SuppressWarnings("unchecked") Coder keyCoder = ((KvCoder) coder).getKeyCoder(); diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ImmutabilityCheckingBundleFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ImmutabilityCheckingBundleFactory.java index 370829451d36..1f157975ecee 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ImmutabilityCheckingBundleFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ImmutabilityCheckingBundleFactory.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; import org.apache.beam.runners.direct.DirectRunner.Enforcement; @@ -31,6 +32,7 @@ import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.HashMultimap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.SetMultimap; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -43,9 +45,6 @@ *

This catches errors during the execution of a {@link DoFn} caused by modifying an element * after it is added to an output {@link PCollection}. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) class ImmutabilityCheckingBundleFactory implements BundleFactory { /** * Create a new {@link ImmutabilityCheckingBundleFactory} that uses the underlying {@link @@ -100,11 +99,12 @@ private class ImmutabilityEnforcingBundle implements UncommittedBundle { public ImmutabilityEnforcingBundle(UncommittedBundle underlying) { this.underlying = underlying; mutationDetectors = HashMultimap.create(); - coder = getPCollection().getCoder(); + // Immutability enforcement only wraps non-root bundles, whose PCollection is always present. + coder = checkStateNotNull(underlying.getPCollection()).getCoder(); } @Override - public PCollection getPCollection() { + public @Nullable PCollection getPCollection() { return underlying.getPCollection(); } @@ -130,7 +130,7 @@ public CommittedBundle commit(Instant synchronizedProcessingTime) { String.format( "PTransform %s mutated value %s after it was output (new value was %s)." + " Values must not be mutated in any way after being output.", - graph.getProducer(underlying.getPCollection()).getFullName(), + graph.getProducer(checkStateNotNull(underlying.getPCollection())).getFullName(), exn.getSavedValue(), exn.getNewValue()), exn.getSavedValue(), diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ImmutabilityEnforcementFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ImmutabilityEnforcementFactory.java index 84509bded752..ca19051f6116 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ImmutabilityEnforcementFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ImmutabilityEnforcementFactory.java @@ -17,6 +17,8 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import java.util.IdentityHashMap; import org.apache.beam.sdk.Pipeline.PipelineVisitor; import org.apache.beam.sdk.coders.Coder; @@ -36,8 +38,7 @@ * element. */ @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) class ImmutabilityEnforcementFactory implements ModelEnforcementFactory { public static ModelEnforcementFactory create() { @@ -112,7 +113,7 @@ private static class ImmutabilityCheckingEnforcement extends AbstractModelEnf private ImmutabilityCheckingEnforcement( CommittedBundle input, AppliedPTransform transform) { this.transform = transform; - coder = input.getPCollection().getCoder(); + coder = checkStateNotNull(input.getPCollection()).getCoder(); mutationElements = new IdentityHashMap<>(); } @@ -128,7 +129,7 @@ public void beforeElement(WindowedValue element) { @Override public void afterElement(WindowedValue element) { - verifyUnmodified(mutationElements.get(element)); + verifyUnmodified(checkStateNotNull(mutationElements.get(element))); } @Override diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ImmutableListBundleFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ImmutableListBundleFactory.java index 5cfb6fa59374..d9f330ea6f48 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ImmutableListBundleFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ImmutableListBundleFactory.java @@ -33,9 +33,6 @@ import org.joda.time.Instant; /** A factory that produces bundles that perform no additional validation. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) class ImmutableListBundleFactory implements BundleFactory { private static final ImmutableListBundleFactory FACTORY = new ImmutableListBundleFactory(); @@ -63,7 +60,7 @@ public UncommittedBundle createKeyedBundle( /** A {@link UncommittedBundle} that buffers elements in memory. */ private static final class UncommittedImmutableListBundle implements UncommittedBundle { - private final PCollection pcollection; + private final @Nullable PCollection pcollection; private final StructuralKey key; private boolean committed = false; private ImmutableList.Builder> elements; @@ -73,18 +70,19 @@ private static final class UncommittedImmutableListBundle implements Uncommit * Create a new {@link UncommittedImmutableListBundle} for the specified {@link PCollection}. */ public static UncommittedImmutableListBundle create( - PCollection pcollection, StructuralKey key) { + @Nullable PCollection pcollection, StructuralKey key) { return new UncommittedImmutableListBundle<>(pcollection, key); } - private UncommittedImmutableListBundle(PCollection pcollection, StructuralKey key) { + private UncommittedImmutableListBundle( + @Nullable PCollection pcollection, StructuralKey key) { this.pcollection = pcollection; this.key = key; this.elements = ImmutableList.builder(); } @Override - public PCollection getPCollection() { + public @Nullable PCollection getPCollection() { return pcollection; } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java index 38938b3ff697..353eb29858a1 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import java.util.HashSet; @@ -46,8 +47,7 @@ // TODO: Handle Key-preserving transforms when appropriate and more aggressively make PTransforms // unkeyed @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) class KeyedPValueTrackingVisitor extends PipelineVisitor.Defaults { @@ -87,7 +87,7 @@ public void leaveCompositeTransform(TransformHierarchy.Node node) { node); if (node.isRootNode()) { finalized = true; - } else if (PRODUCES_KEYED_OUTPUTS.contains(node.getTransform().getClass())) { + } else if (PRODUCES_KEYED_OUTPUTS.contains(checkStateNotNull(node.getTransform()).getClass())) { Map, PCollection> outputs = node.getOutputs(); for (PValue output : outputs.values()) { keyedValues.add(output); @@ -101,8 +101,9 @@ public void visitValue(PValue value, TransformHierarchy.Node producer) { for (PValue input : producer.getInputs().values()) { inputsAreKeyed = inputsAreKeyed && keyedValues.contains(input); } - if (PRODUCES_KEYED_OUTPUTS.contains(producer.getTransform().getClass()) - || (isKeyPreserving(producer.getTransform()) && inputsAreKeyed)) { + PTransform transform = checkStateNotNull(producer.getTransform()); + if (PRODUCES_KEYED_OUTPUTS.contains(transform.getClass()) + || (isKeyPreserving(transform) && inputsAreKeyed)) { keyedValues.add(value); } } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/MultiStepCombine.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/MultiStepCombine.java index bd8613da6118..5ec3ebd9e123 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/MultiStepCombine.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/MultiStepCombine.java @@ -65,7 +65,11 @@ /** A {@link Combine} that performs the combine in multiple steps. */ @SuppressWarnings({ "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + // The "nullness" suppression is retained deliberately: this class is generic over a nullable + // AccumT (e.g. an empty HllCount sketch has a null accumulator), so accumulator values, keys and + // combine inputs are legitimately nullable in normal operation. The checker cannot express that + // invariant here, and forcing @NonNull bounds converts valid nulls into runtime failures. + "nullness" }) class MultiStepCombine< K extends @Nullable Object, diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluator.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluator.java index 6124f4a5413d..8fb582ab04ed 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluator.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluator.java @@ -17,7 +17,7 @@ */ package org.apache.beam.runners.direct; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; +import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; import java.util.HashMap; import java.util.List; @@ -48,10 +48,10 @@ import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) class ParDoEvaluator implements TransformEvaluator { @@ -272,7 +272,7 @@ public TransformResult finishBundle() { throw UserCodeException.wrap(e); } StepTransformResult.Builder resultBuilder; - CopyOnAccessInMemoryStateInternals state = stepContext.commitState(); + @Nullable CopyOnAccessInMemoryStateInternals state = stepContext.commitState(); if (state != null) { resultBuilder = StepTransformResult.withHold(transform, state.getEarliestWatermarkHold()) @@ -301,8 +301,9 @@ private BundleOutputManager(Map, UncommittedBundle> bundles) { @Override public void output(TupleTag tag, WindowedValue output) { - checkArgument(bundles.containsKey(tag), "Unknown output tag %s", tag); - bundles.get(tag).add((WindowedValue) output); + UncommittedBundle bundle = + checkArgumentNotNull(bundles.get(tag), "Unknown output tag %s", tag); + bundle.add((WindowedValue) output); } } } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluatorFactory.java index 22729d91720a..c545deb38832 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluatorFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluatorFactory.java @@ -17,6 +17,8 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import java.util.List; import java.util.Map; import org.apache.beam.runners.direct.DirectExecutionContext.DirectStepContext; @@ -39,9 +41,6 @@ import org.slf4j.LoggerFactory; /** A {@link TransformEvaluatorFactory} for {@link ParDo.MultiOutput}. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) final class ParDoEvaluatorFactory implements TransformEvaluatorFactory { private static final Logger LOG = LoggerFactory.getLogger(ParDoEvaluatorFactory.class); @@ -80,7 +79,7 @@ public TransformEvaluator forApplication( (TransformEvaluator) createEvaluator( (AppliedPTransform) application, - (PCollection) inputBundle.getPCollection(), + (PCollection) checkStateNotNull(inputBundle.getPCollection()), inputBundle.getKey(), ParDoTranslation.getSideInputs(application), (TupleTag) ParDoTranslation.getMainOutputTag(application), diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/PipelineExecutor.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/PipelineExecutor.java index e851541b44b9..e52385c30578 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/PipelineExecutor.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/PipelineExecutor.java @@ -21,6 +21,7 @@ import org.apache.beam.sdk.PipelineResult.State; import org.apache.beam.sdk.runners.AppliedPTransform; import org.apache.beam.sdk.transforms.PTransform; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; /** @@ -46,7 +47,7 @@ interface PipelineExecutor { * @throws Exception whenever an executor thread throws anything, transfers to the waiting thread * and rethrows it */ - State waitUntilFinish(Duration duration) throws Exception; + @Nullable State waitUntilFinish(Duration duration) throws Exception; /** Gets the current state of the {@link Pipeline}. */ State getPipelineState(); diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/QuiescenceDriver.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/QuiescenceDriver.java index 8915c25bef87..78d5d3180815 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/QuiescenceDriver.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/QuiescenceDriver.java @@ -17,6 +17,8 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import com.google.auto.value.AutoValue; import java.util.ArrayList; import java.util.Collection; @@ -48,9 +50,8 @@ */ @SuppressWarnings({ "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "keyfor", - "nullness" -}) // TODO(https://github.com/apache/beam/issues/20497) + "keyfor" +}) class QuiescenceDriver implements ExecutionDriver { private static final Logger LOG = LoggerFactory.getLogger(QuiescenceDriver.class); @@ -228,7 +229,7 @@ private void addWorkIfNecessary() { // Pull all available work off of the queue, then schedule it all, so this loop // terminates while (!pendingRootEntry.getValue().isEmpty()) { - CommittedBundle bundle = pendingRootEntry.getValue().poll(); + CommittedBundle bundle = checkStateNotNull(pendingRootEntry.getValue().poll()); bundles.add(bundle); } for (CommittedBundle bundle : bundles) { @@ -299,7 +300,8 @@ public final CommittedResult handleResult( for (CommittedBundle outputBundle : committedResult.getOutputs()) { pendingWork.offer( WorkUpdate.fromBundle( - outputBundle, graph.getPerElementConsumers(outputBundle.getPCollection()))); + outputBundle, + graph.getPerElementConsumers(checkStateNotNull(outputBundle.getPCollection())))); } Optional> unprocessedInputs = committedResult.getUnprocessedInputs(); @@ -307,7 +309,8 @@ public final CommittedResult handleResult( if (inputBundle.getPCollection() == null) { // TODO: Split this logic out of an if statement synchronized (pendingRootBundles) { - pendingRootBundles.get(result.getTransform()).offer(unprocessedInputs.get()); + checkStateNotNull(pendingRootBundles.get(result.getTransform())) + .offer(unprocessedInputs.get()); } } else { pendingWork.offer( @@ -319,12 +322,12 @@ public final CommittedResult handleResult( state.set(ExecutorState.ACTIVE); } synchronized (inflightBundles) { - inflightBundles.compute( - result.getTransform(), - (k, v) -> { - v.remove(inputBundle); - return v.isEmpty() ? null : v; - }); + Collection> remaining = + checkStateNotNull(inflightBundles.get(result.getTransform())); + remaining.remove(inputBundle); + if (remaining.isEmpty()) { + inflightBundles.remove(result.getTransform()); + } } outstandingWork.decrementAndGet(); return committedResult; diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/RootProviderRegistry.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/RootProviderRegistry.java index 7ac36e2fa8b8..527d92a514e0 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/RootProviderRegistry.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/RootProviderRegistry.java @@ -36,8 +36,7 @@ * based on the type of {@link PTransform} of the application. */ @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) class RootProviderRegistry { /** Returns a {@link RootProviderRegistry} that supports the Java SDK root transforms. */ diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/SideInputContainer.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/SideInputContainer.java index acbce34e0f95..276aa0781af4 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/SideInputContainer.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/SideInputContainer.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import java.util.ArrayList; @@ -57,8 +58,7 @@ * and writing to a {@link PCollectionView}. */ @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) class SideInputContainer { private static final Set SUPPORTED_MATERIALIZATIONS = @@ -68,7 +68,7 @@ class SideInputContainer { private final Collection> containedViews; private final LoadingCache< - PCollectionViewWindow, AtomicReference>>> + PCollectionViewWindow, AtomicReference<@Nullable Iterable>>> viewByWindows; /** Create a new {@link SideInputContainer} with the provided views and the provided context. */ @@ -84,14 +84,18 @@ public static SideInputContainer create( pCollectionView.getViewFn().getMaterialization().getUrn(), pCollectionView.getTagInternal().getId()); } - LoadingCache, AtomicReference>>> + LoadingCache< + PCollectionViewWindow, + AtomicReference<@Nullable Iterable>>> viewByWindows = CacheBuilder.newBuilder().build(new CallbackSchedulingLoader(context)); return new SideInputContainer(containedViews, viewByWindows); } private SideInputContainer( Collection> containedViews, - LoadingCache, AtomicReference>>> + LoadingCache< + PCollectionViewWindow, + AtomicReference<@Nullable Iterable>>> viewByWindows) { this.containedViews = ImmutableSet.copyOf(containedViews); this.viewByWindows = viewByWindows; @@ -153,7 +157,7 @@ private Map>> indexValuesByWindow( private void updatePCollectionViewWindowValues( PCollectionView view, BoundedWindow window, Collection> windowValues) { PCollectionViewWindow windowedView = PCollectionViewWindow.of(view, window); - AtomicReference>> contents = + AtomicReference<@Nullable Iterable>> contents = viewByWindows.getUnchecked(windowedView); if (contents.compareAndSet(null, windowValues)) { // the value had never been set, so we set it and are done. @@ -164,7 +168,7 @@ private void updatePCollectionViewWindowValues( Iterable> existingValues; long existingPane; do { - existingValues = contents.get(); + existingValues = checkStateNotNull(contents.get()); existingPane = Iterables.isEmpty(existingValues) ? -1L @@ -175,7 +179,8 @@ private void updatePCollectionViewWindowValues( private static class CallbackSchedulingLoader extends CacheLoader< - PCollectionViewWindow, AtomicReference>>> { + PCollectionViewWindow, + AtomicReference<@Nullable Iterable>>> { private final EvaluationContext context; public CallbackSchedulingLoader(EvaluationContext context) { @@ -183,10 +188,11 @@ public CallbackSchedulingLoader(EvaluationContext context) { } @Override - public AtomicReference>> load( + public AtomicReference<@Nullable Iterable>> load( PCollectionViewWindow view) { - AtomicReference>> contents = new AtomicReference<>(); + AtomicReference<@Nullable Iterable>> contents = + new AtomicReference<>(); WindowingStrategy windowingStrategy = view.getView().getWindowingStrategyInternal(); context.scheduleAfterOutputWouldBeProduced( @@ -201,12 +207,12 @@ public AtomicReference>> load( private static class WriteEmptyViewContents implements Runnable { private final PCollectionView view; private final BoundedWindow window; - private final AtomicReference>> contents; + private final AtomicReference<@Nullable Iterable>> contents; private WriteEmptyViewContents( PCollectionView view, BoundedWindow window, - AtomicReference>> contents) { + AtomicReference<@Nullable Iterable>> contents) { this.contents = contents; this.view = view; this.window = window; diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/SplittableProcessElementsEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/SplittableProcessElementsEvaluatorFactory.java index b134e872b65d..afd12a4a6cfa 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/SplittableProcessElementsEvaluatorFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/SplittableProcessElementsEvaluatorFactory.java @@ -17,10 +17,12 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.atomic.AtomicInteger; import org.apache.beam.runners.core.DoFnRunners; import org.apache.beam.runners.core.KeyedWorkItem; import org.apache.beam.runners.core.OutputAndTimeBoundedSplittableProcessElementInvoker; @@ -39,12 +41,11 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder; import org.joda.time.Duration; -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) class SplittableProcessElementsEvaluatorFactory< InputT, OutputT, RestrictionT, PositionT, WatermarkEstimatorStateT> implements TransformEvaluatorFactory { + private static final AtomicInteger FACTORY_INSTANCE_IDS = new AtomicInteger(); + private final ParDoEvaluatorFactory>, OutputT> delegateFactory; private final ScheduledExecutorService ses; @@ -82,7 +83,8 @@ public DoFnLifecycleManager load(final AppliedPTransform application) { new ThreadFactoryBuilder() .setThreadFactory(MoreExecutors.platformThreadFactory()) .setNameFormat( - "direct-splittable-process-element-checkpoint-executor_" + hashCode()) + "direct-splittable-process-element-checkpoint-executor_" + + FACTORY_INSTANCE_IDS.getAndIncrement()) .build()); } @@ -120,7 +122,7 @@ private TransformEvaluator>> crea delegateFactory.createEvaluator( (AppliedPTransform) application, (PCollection>>) - inputBundle.getPCollection(), + checkStateNotNull(inputBundle.getPCollection()), inputBundle.getKey(), application.getTransform().getSideInputs(), application.getTransform().getMainOutputTag(), diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactory.java index 8edbd45f2b5d..1e8b6ccc4116 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactory.java @@ -17,6 +17,8 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import com.google.auto.value.AutoValue; import java.util.Collections; import java.util.NavigableSet; @@ -47,8 +49,7 @@ /** A {@link TransformEvaluatorFactory} for stateful {@link ParDo}. */ @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) final class StatefulParDoEvaluatorFactory implements TransformEvaluatorFactory { @@ -104,7 +105,7 @@ private TransformEvaluator>> createEvaluator( DoFnLifecycleManagerRemovingTransformEvaluator> delegateEvaluator = delegateFactory.createEvaluator( (AppliedPTransform) application, - (PCollection) inputBundle.getPCollection(), + (PCollection) checkStateNotNull(inputBundle.getPCollection()), inputBundle.getKey(), application.getTransform().getSideInputs(), application.getTransform().getMainOutputTag(), @@ -176,7 +177,7 @@ public void processElement(WindowedValue>> gbkRes NavigableSet earlierTimers = timerInternals.getModifiedTimersOrdered(timerData.getDomain()).headSet(timerData, true); while (!earlierTimers.isEmpty()) { - TimerData insertedTimer = earlierTimers.pollFirst(); + TimerData insertedTimer = checkStateNotNull(earlierTimers.pollFirst()); if (timerModified(insertedTimer)) { continue; } @@ -237,17 +238,21 @@ public TransformResult>> finishBundle() throws Ex clearWatermarkHold(timerData); } - CopyOnAccessInMemoryStateInternals state; + @Nullable CopyOnAccessInMemoryStateInternals delegateState = delegateResult.getState(); + @Nullable CopyOnAccessInMemoryStateInternals state; Instant watermarkHold; - if (isTimerDeclared && delegateResult.getState() != null) { // For both State and Timer Holds - state = delegateResult.getState(); - watermarkHold = stepContext.commitState().getEarliestWatermarkHold(); + if (isTimerDeclared && delegateState != null) { // For both State and Timer Holds + state = delegateState; + watermarkHold = checkStateNotNull(stepContext.commitState()).getEarliestWatermarkHold(); } else if (isTimerDeclared) { // For only Timer holds - state = stepContext.commitState(); - watermarkHold = state.getEarliestWatermarkHold(); + // A declared timer accessed state via setWatermarkHold, so commitState is non-null. + CopyOnAccessInMemoryStateInternals committedState = + checkStateNotNull(stepContext.commitState()); + state = committedState; + watermarkHold = committedState.getEarliestWatermarkHold(); } else { // For only State ( non Timer ) holds - state = delegateResult.getState(); + state = delegateState; watermarkHold = delegateResult.getWatermarkHold(); } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StepTransformResult.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StepTransformResult.java index 8cfe07bdbb71..c9ec6ec64643 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StepTransformResult.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StepTransformResult.java @@ -33,14 +33,14 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** An immutable {@link TransformResult}. */ @AutoValue @AutoValue.CopyAnnotations @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) abstract class StepTransformResult implements TransformResult { @@ -73,7 +73,7 @@ public static class Builder { private final ImmutableList.Builder> bundlesBuilder; private final ImmutableList.Builder> unprocessedElementsBuilder; private MetricUpdates metricUpdates; - private CopyOnAccessInMemoryStateInternals state; + private @Nullable CopyOnAccessInMemoryStateInternals state; private TimerUpdate timerUpdate; private List finalizations; private final Set producedOutputs; @@ -85,7 +85,7 @@ private Builder(AppliedPTransform transform, Instant watermarkHold) { this.bundlesBuilder = ImmutableList.builder(); this.producedOutputs = EnumSet.noneOf(OutputType.class); this.unprocessedElementsBuilder = ImmutableList.builder(); - this.timerUpdate = TimerUpdate.builder(null).build(); + this.timerUpdate = TimerUpdate.empty(); this.metricUpdates = MetricUpdates.EMPTY; this.finalizations = Collections.EMPTY_LIST; } @@ -108,7 +108,7 @@ public Builder withMetricUpdates(MetricUpdates metricUpdates) { return this; } - public Builder withState(CopyOnAccessInMemoryStateInternals state) { + public Builder withState(@Nullable CopyOnAccessInMemoryStateInternals state) { this.state = state; return this; } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactory.java index 8ca2a909e367..c1bc06f069cd 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactory.java @@ -52,8 +52,7 @@ /** The {@link TransformEvaluatorFactory} for the {@link TestStream} primitive. */ @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) class TestStreamEvaluatorFactory implements TransformEvaluatorFactory { private final EvaluationContext evaluationContext; diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformEvaluatorRegistry.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformEvaluatorRegistry.java index 544222c8253f..e416f85f0d78 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformEvaluatorRegistry.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformEvaluatorRegistry.java @@ -23,6 +23,7 @@ import static org.apache.beam.runners.direct.MultiStepCombine.DIRECT_MERGE_ACCUMULATORS_EXTRACT_OUTPUT_URN; import static org.apache.beam.runners.direct.ParDoMultiOverrideFactory.DIRECT_STATEFUL_PAR_DO_URN; import static org.apache.beam.runners.direct.TestStreamEvaluatorFactory.DirectTestStreamFactory.DIRECT_TEST_STREAM_URN; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.sdk.util.construction.PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN; import static org.apache.beam.sdk.util.construction.PTransformTranslation.CREATE_VIEW_TRANSFORM_URN; import static org.apache.beam.sdk.util.construction.PTransformTranslation.FLATTEN_TRANSFORM_URN; @@ -56,8 +57,7 @@ * implementations based on the type of {@link PTransform} of the application. */ @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) class TransformEvaluatorRegistry { private static final Logger LOG = LoggerFactory.getLogger(TransformEvaluatorRegistry.class); @@ -181,7 +181,7 @@ public void cleanup() throws Exception { toThrow.addSuppressed(e); } } - throw toThrow; + throw checkStateNotNull(toThrow); } } } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformExecutorServices.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformExecutorServices.java index 1ec5b1c3f17e..eab5e7534658 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformExecutorServices.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformExecutorServices.java @@ -24,13 +24,11 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; +import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** Static factory methods for constructing instances of {@link TransformExecutorService}. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) final class TransformExecutorServices { private TransformExecutorServices() { // Do not instantiate @@ -115,7 +113,7 @@ public void shutdown() { private static class SerialTransformExecutor implements TransformExecutorService { private final ExecutorService executor; - private AtomicReference currentlyEvaluating; + private AtomicReference<@Nullable TransformExecutor> currentlyEvaluating; private final Queue workQueue; private boolean active = true; diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/UnboundedReadEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/UnboundedReadEvaluatorFactory.java index c4a4fe86bf90..29b612e05399 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/UnboundedReadEvaluatorFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/UnboundedReadEvaluatorFactory.java @@ -52,8 +52,7 @@ * PTransform}. */ @SuppressWarnings({ - "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) class UnboundedReadEvaluatorFactory implements TransformEvaluatorFactory { // Occasionally close an existing reader and resume from checkpoint, to exercise close-and-resume @@ -239,10 +238,11 @@ private UnboundedReader getReader(UnboundedSourceShard reader, UnboundedSourceShard shard) throws IOException { - if (shard.getExistingReader() == null) { + UnboundedReader existingReader = shard.getExistingReader(); + if (existingReader == null) { return reader.start(); } else { - return shard.getExistingReader().advance(); + return existingReader.advance(); } } @@ -298,6 +298,7 @@ static UnboundedSourceShard UnboundedSourceShard of( UnboundedSource source, UnboundedReadDeduplicator deduplicator, diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkCallbackExecutor.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkCallbackExecutor.java index 9bfbb67fad2c..a66683f3dee8 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkCallbackExecutor.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkCallbackExecutor.java @@ -17,6 +17,8 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import java.io.Serializable; import java.util.ArrayList; @@ -51,9 +53,6 @@ * followed by a call to {@link #fireForWatermark(AppliedPTransform, Instant)} for the same * transform with the current value of the watermark. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) class WatermarkCallbackExecutor { /** Create a new {@link WatermarkCallbackExecutor}. */ public static WatermarkCallbackExecutor create(Executor executor) { @@ -84,10 +83,9 @@ public void callOnGuaranteedFiring( PriorityQueue callbackQueue = callbacks.get(step); if (callbackQueue == null) { - callbackQueue = new PriorityQueue<>(11, new CallbackOrdering()); - if (callbacks.putIfAbsent(step, callbackQueue) != null) { - callbackQueue = callbacks.get(step); - } + PriorityQueue newQueue = new PriorityQueue<>(11, new CallbackOrdering()); + PriorityQueue existing = callbacks.putIfAbsent(step, newQueue); + callbackQueue = existing != null ? existing : newQueue; } synchronized (callbackQueue) { @@ -109,10 +107,9 @@ public void callOnWindowExpiration( PriorityQueue callbackQueue = callbacks.get(step); if (callbackQueue == null) { - callbackQueue = new PriorityQueue<>(11, new CallbackOrdering()); - if (callbacks.putIfAbsent(step, callbackQueue) != null) { - callbackQueue = callbacks.get(step); - } + PriorityQueue newQueue = new PriorityQueue<>(11, new CallbackOrdering()); + PriorityQueue existing = callbacks.putIfAbsent(step, newQueue); + callbackQueue = existing != null ? existing : newQueue; } synchronized (callbackQueue) { @@ -132,8 +129,9 @@ public void fireForWatermark(AppliedPTransform step, Instant watermark) } synchronized (callbackQueue) { List toFire = new ArrayList<>(); - while (!callbackQueue.isEmpty() && callbackQueue.peek().shouldFire(watermark)) { - toFire.add(callbackQueue.poll().getCallback()); + while (!callbackQueue.isEmpty() + && checkStateNotNull(callbackQueue.peek()).shouldFire(watermark)) { + toFire.add(checkStateNotNull(callbackQueue.poll()).getCallback()); } if (!toFire.isEmpty()) { CountDownLatch latch = new CountDownLatch(toFire.size()); diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java index b048b4e82c18..47530da5f6cc 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.direct; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; @@ -66,6 +67,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.SortedMultiset; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.TreeMultiset; +import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; @@ -129,10 +131,7 @@ * Watermark_PCollection = Watermark_Out_ProducingPTransform * */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) -class WatermarkManager { +class WatermarkManager { // The number of updates to apply in #tryApplyPendingUpdates private static final int MAX_INCREMENTAL_UPDATES = 10; @@ -233,11 +232,11 @@ static class AppliedPTransformInputWatermark implements Watermark { private final SortedMultiset pendingTimers; // Entries in this table represent the authoritative timestamp for which - // a per-key-and-StateNamespace timer is set. - private final Map, Map> existingTimers; + // a per-key-and-StateNamespace timer is set. The key is null for a keyless (empty) TimerUpdate. + private final Map<@Nullable StructuralKey, Map> existingTimers; // This per-key sorted set allows quick retrieval of timers that should fire for a key - private final Map, NavigableSet> objectTimers; + private final Map<@Nullable StructuralKey, NavigableSet> objectTimers; private final AtomicReference currentWatermark; @@ -299,9 +298,11 @@ public synchronized WatermarkUpdate refresh() { minInputWatermark = INSTANT_ORDERING.min(minInputWatermark, inputWatermark.get()); } if (!pendingElements.isEmpty()) { + // Guarded by the isEmpty() check above, so firstEntry() is always present. minInputWatermark = INSTANT_ORDERING.min( - minInputWatermark, pendingElements.firstEntry().getElement().getMinimumTimestamp()); + minInputWatermark, + checkStateNotNull(pendingElements.firstEntry()).getElement().getMinimumTimestamp()); } Instant newWatermark = INSTANT_ORDERING.max(oldWatermark, minInputWatermark); currentWatermark.set(newWatermark); @@ -326,7 +327,9 @@ synchronized Instant getEarliestTimerTimestamp() { } private Instant getMinimumOutputTimestamp(SortedMultiset timers) { - Instant minimumOutputTimestamp = timers.firstEntry().getElement().getOutputTimestamp(); + // Only called with a non-empty multiset, so firstEntry() is always present. + Instant minimumOutputTimestamp = + checkStateNotNull(timers.firstEntry()).getElement().getOutputTimestamp(); for (TimerData timerData : timers) { minimumOutputTimestamp = INSTANT_ORDERING.min(timerData.getOutputTimestamp(), minimumOutputTimestamp); @@ -430,7 +433,7 @@ public AppliedPTransformOutputWatermark( currentWatermark = new AtomicReference<>(BoundedWindow.TIMESTAMP_MIN_VALUE); } - public synchronized void updateHold(Object key, Instant newHold) { + public synchronized void updateHold(@Nullable Object key, Instant newHold) { if (newHold == null) { holds.removeHold(key); } else { @@ -509,7 +512,8 @@ private static class SynchronizedProcessingTimeInputWatermark implements Waterma private final Collection> pendingBundles; private final Map, NavigableSet> processingTimers; private final Map, NavigableSet> synchronizedProcessingTimers; - private final Map, Map> existingTimers; + // The key is null for a keyless (empty) TimerUpdate; keyed timers always carry a non-null key. + private final Map<@Nullable StructuralKey, Map> existingTimers; private final NavigableSet pendingTimers; @@ -622,10 +626,12 @@ private synchronized void updateTimers(TimerUpdate update) { Map existingTimersForKey = existingTimers.computeIfAbsent(update.key, k -> Maps.newHashMap()); + // A TimerUpdate carrying set or deleted timers was built via builder(key), so its key is + // non-null; only the keyless empty() update has a null key and it carries no such timers. HashSet newSetTimers = Sets.newHashSet(); for (TimerData addedTimer : update.setTimers.values()) { NavigableSet timerQueue = - processQueueForDomain(update.key, addedTimer.getDomain()); + processQueueForDomain(checkStateNotNull(update.key), addedTimer.getDomain()); if (timerQueue == null) { continue; } @@ -643,7 +649,7 @@ private synchronized void updateTimers(TimerUpdate update) { for (TimerData deletedTimer : update.deletedTimers.values()) { NavigableSet timerQueue = - processQueueForDomain(update.key, deletedTimer.getDomain()); + processQueueForDomain(checkStateNotNull(update.key), deletedTimer.getDomain()); if (timerQueue == null) { continue; } @@ -749,7 +755,7 @@ public SynchronizedProcessingTimeOutputWatermark( this.latestRefresh = new AtomicReference<>(BoundedWindow.TIMESTAMP_MIN_VALUE); } - public synchronized void updateHold(Object key, Instant newHold) { + public synchronized void updateHold(@Nullable Object key, Instant newHold) { if (newHold == null) { holds.removeHold(key); } else { @@ -838,15 +844,17 @@ public Instant get() { *

The result collection retains ordering of timers (from earliest to latest). */ private static synchronized Map, List> extractFiredTimers( - Instant latestTime, Map, NavigableSet> objectTimers) { + Instant latestTime, + Map, NavigableSet> objectTimers) { Map, List> result = new HashMap<>(); - Set> emptyKeys = new HashSet<>(); - for (Map.Entry, NavigableSet> pendingTimers : + Set<@Nullable StructuralKey> emptyKeys = new HashSet<>(); + for (Map.Entry, NavigableSet> pendingTimers : objectTimers.entrySet()) { NavigableSet timers = pendingTimers.getValue(); if (!timers.isEmpty() && timers.first().getTimestamp().isBefore(latestTime)) { ArrayList keyFiredTimers = new ArrayList<>(); - result.put(pendingTimers.getKey(), keyFiredTimers); + // A non-empty timer set is only ever stored under a non-null key. + result.put(checkStateNotNull(pendingTimers.getKey()), keyFiredTimers); while (!timers.isEmpty() && timers.first().getTimestamp().isBefore(latestTime)) { keyFiredTimers.add(timers.first()); timers.remove(timers.first()); @@ -909,12 +917,17 @@ private static synchronized Map, List> extractFiredT * @param getName a function for producing a short identifier for the executable in watermark * tracing log messages. */ - public static - WatermarkManager create( + public static + WatermarkManager create( Clock clock, - ExecutableGraph graph, + ExecutableGraph graph, Function getName) { - return new WatermarkManager<>(clock, graph, getName); + WatermarkManager watermarkManager = + new WatermarkManager<>(clock, graph, getName); + // Populate the per-transform watermarks after construction completes, because + // getTransformWatermark reads instance state and must run on a fully-initialized receiver. + watermarkManager.initializeTransformWatermarks(); + return watermarkManager; } private WatermarkManager( @@ -931,7 +944,9 @@ private WatermarkManager( this.pendingRefreshes = new HashSet<>(); transformToWatermarks = new HashMap<>(); + } + private void initializeTransformWatermarks() { for (ExecutableT rootTransform : graph.getRootTransforms()) { getTransformWatermark(rootTransform); } @@ -979,23 +994,19 @@ private TransformWatermarks getTransformWatermark(ExecutableT executable) { return wms; } - private static Consumer timerUpdateConsumer( + private static Consumer timerUpdateConsumer( Map> transformsWithAlreadyExtractedTimers, ExecutableT executable) { return update -> { String timerIdWithNs = update.stringKey(); synchronized (transformsWithAlreadyExtractedTimers) { - transformsWithAlreadyExtractedTimers.compute( - executable, - (k, v) -> { - if (v != null) { - v.remove(timerIdWithNs); - if (v.isEmpty()) { - v = null; - } - } - return v; - }); + Set extracted = transformsWithAlreadyExtractedTimers.get(executable); + if (extracted != null) { + extracted.remove(timerIdWithNs); + if (extracted.isEmpty()) { + transformsWithAlreadyExtractedTimers.remove(executable); + } + } } }; } @@ -1037,7 +1048,8 @@ private List getInputWatermarks(ExecutableT executable) { * @return a snapshot of the input watermark and output watermark for the provided executable */ public TransformWatermarks getWatermarks(ExecutableT executable) { - return transformToWatermarks.get(executable); + // Every executable in the graph is registered during construction, so a watermark is present. + return checkStateNotNull(transformToWatermarks.get(executable)); } public void initialize( @@ -1046,7 +1058,8 @@ public void initialize( try { for (Map.Entry>> rootEntry : initialBundles.entrySet()) { - TransformWatermarks rootWms = transformToWatermarks.get(rootEntry.getKey()); + TransformWatermarks rootWms = + checkStateNotNull(transformToWatermarks.get(rootEntry.getKey())); for (Bundle initialBundle : rootEntry.getValue()) { rootWms.addPending(initialBundle); } @@ -1129,7 +1142,9 @@ private void applyAllPendingUpdates() { private void applyNUpdates(int numUpdates) { synchronized (pendingUpdates) { for (int i = 0; !pendingUpdates.isEmpty() && ((i < numUpdates) || (numUpdates <= 0)); i++) { - PendingWatermarkUpdate pending = pendingUpdates.poll(); + // Guarded by the isEmpty() check in the loop condition, so poll() returns non-null. + PendingWatermarkUpdate pending = + checkStateNotNull(pendingUpdates.poll()); applyPendingUpdate(pending); pendingRefreshes.add(pending.getExecutable()); } @@ -1148,7 +1163,7 @@ private void applyPendingUpdate(PendingWatermarkUpdate pending.getUnprocessedInputs(), pending.getOutputs()); - TransformWatermarks transformWms = transformToWatermarks.get(executable); + TransformWatermarks transformWms = checkStateNotNull(transformToWatermarks.get(executable)); transformWms.setEventTimeHold( inputBundle == null ? null : inputBundle.getKey(), pending.getEarliestHold()); @@ -1172,7 +1187,7 @@ private void applyPendingUpdate(PendingWatermarkUpdate * information about the parameters of this method. */ private void updatePending( - Bundle input, + @Nullable Bundle input, TimerUpdate timerUpdate, ExecutableT executable, @Nullable Bundle unprocessedInputs, @@ -1181,13 +1196,15 @@ private void updatePending( // do not share a Mutex within this call and thus can be interleaved with external calls to // refresh. for (Bundle bundle : outputs) { - for (ExecutableT consumer : graph.getPerElementConsumers(bundle.getPCollection())) { - TransformWatermarks watermarks = transformToWatermarks.get(consumer); + for (ExecutableT consumer : + graph.getPerElementConsumers(checkStateNotNull(bundle.getPCollection()))) { + TransformWatermarks watermarks = checkStateNotNull(transformToWatermarks.get(consumer)); watermarks.addPending(bundle); } } - TransformWatermarks completedTransform = transformToWatermarks.get(executable); + TransformWatermarks completedTransform = + checkStateNotNull(transformToWatermarks.get(executable)); if (unprocessedInputs != null) { // Add the unprocessed inputs completedTransform.addPending(unprocessedInputs); @@ -1225,7 +1242,7 @@ private Set refreshAllOf(Set toRefresh) { } private Set refreshWatermarks(final ExecutableT toRefresh) { - TransformWatermarks myWatermarks = transformToWatermarks.get(toRefresh); + TransformWatermarks myWatermarks = checkStateNotNull(transformToWatermarks.get(toRefresh)); WatermarkUpdate updateResult = myWatermarks.refresh(); if (updateResult.isAdvanced()) { Set additionalRefreshes = new HashSet<>(); @@ -1296,17 +1313,17 @@ public Collection> extractFiredTimers( * as the key is arbitrarily ordered via identity, rather than object equality. */ private static final class KeyedHold implements Comparable { - private static final Ordering KEY_ORDERING = Ordering.arbitrary().nullsLast(); + private static final Ordering<@Nullable Object> KEY_ORDERING = Ordering.arbitrary().nullsLast(); - private final Object key; + private final @Nullable Object key; private final Instant timestamp; /** Create a new KeyedHold with the specified key and timestamp. */ - public static KeyedHold of(Object key, Instant timestamp) { + public static KeyedHold of(@Nullable Object key, Instant timestamp) { return new KeyedHold(key, MoreObjects.firstNonNull(timestamp, THE_END_OF_TIME.get())); } - private KeyedHold(Object key, Instant timestamp) { + private KeyedHold(@Nullable Object key, Instant timestamp) { this.key = key; this.timestamp = timestamp; } @@ -1349,6 +1366,9 @@ public String toString() { private static class PerKeyHolds { private final Map keyedHolds; + // The single hold for a keyless (null-key) update, tracked separately so keyedHolds never needs + // a null map key. + private @Nullable KeyedHold keylessHold; private final NavigableSet allHolds; private PerKeyHolds() { @@ -1371,13 +1391,23 @@ public Instant getMinHold() { public void updateHold(@Nullable Object key, Instant newHold) { removeHold(key); KeyedHold newKeyedHold = KeyedHold.of(key, newHold); - keyedHolds.put(key, newKeyedHold); + if (key == null) { + keylessHold = newKeyedHold; + } else { + keyedHolds.put(key, newKeyedHold); + } allHolds.add(newKeyedHold); } /** Removes the hold of the provided key. */ - public void removeHold(Object key) { - KeyedHold oldHold = keyedHolds.remove(key); + public void removeHold(@Nullable Object key) { + KeyedHold oldHold; + if (key == null) { + oldHold = keylessHold; + keylessHold = null; + } else { + oldHold = keyedHolds.remove(key); + } if (oldHold != null) { allHolds.remove(oldHold); } @@ -1459,11 +1489,11 @@ private WatermarkUpdate refresh() { return eventOutputUpdate.union(syncOutputUpdate); } - private void setEventTimeHold(Object key, Instant newHold) { + private void setEventTimeHold(@Nullable Object key, Instant newHold) { outputWatermark.updateHold(key, newHold); } - private void setSynchronizedProcessingTimeHold(Object key, Instant newHold) { + private void setSynchronizedProcessingTimeHold(@Nullable Object key, Instant newHold) { synchronizedProcessingOutputWatermark.updateHold(key, newHold); } @@ -1556,7 +1586,7 @@ static TimerKey of(TimerData timerData) { * the input to the executed step. */ public static class TimerUpdate { - private final StructuralKey key; + private final @Nullable StructuralKey key; private final Iterable completedTimers; private final Map setTimers; private final Map deletedTimers; @@ -1640,7 +1670,7 @@ public TimerUpdate build() { } private TimerUpdate( - StructuralKey key, + @Nullable StructuralKey key, Iterable completedTimers, Map setTimers, Map deletedTimers) { @@ -1651,7 +1681,7 @@ private TimerUpdate( } @VisibleForTesting - StructuralKey getKey() { + @Nullable StructuralKey getKey() { return key; } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WindowEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WindowEvaluatorFactory.java index 2724312c99a7..b032fe5bbf47 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WindowEvaluatorFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WindowEvaluatorFactory.java @@ -35,7 +35,7 @@ * The {@link DirectRunner} {@link TransformEvaluatorFactory} for the {@link Window.Assign} * primitive {@link PTransform}. */ -@SuppressWarnings({"keyfor", "nullness"}) // TODO(https://github.com/apache/beam/issues/20497) +@SuppressWarnings("keyfor") class WindowEvaluatorFactory implements TransformEvaluatorFactory { private final EvaluationContext evaluationContext;