diff --git a/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto b/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto index 9360522ab409..9b32048b4995 100644 --- a/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto +++ b/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto @@ -740,10 +740,18 @@ message Elements { bool is_last = 4; } + message DrainMode { + enum Enum { + UNSPECIFIED = 0; + NOT_DRAINING = 1; + DRAINING = 2; + } + } + // Element metadata passed as part of WindowedValue to make WindowedValue // extensible and backward compatible message ElementMetadata { - // empty message - add drain, kind, tracing metadata in the future + optional DrainMode.Enum drain = 1; } // Represent the encoded user timer for a given instruction, transform and diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java index dce4fcaa68bd..02b68827d2e0 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java @@ -46,7 +46,8 @@ public interface DoFnRunner void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { - doFnRunner.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + TimeDomain timeDomain, + @Nullable Boolean draining) { + doFnRunner.onTimer( + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } @Override diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java index 9bda4dd2cbca..5203c10c9deb 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java @@ -386,6 +386,11 @@ public PaneInfo pane() { return element.getPaneInfo(); } + @Override + public Boolean draining() { + return element.isDraining(); + } + @Override public String currentRecordId() { return element.getRecordId(); @@ -439,7 +444,13 @@ public void outputWindowedValue( outputReceiver.output( mainOutputTag, WindowedValues.of( - value, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + value, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + element.isDraining())); } @Override @@ -450,7 +461,15 @@ public void output(TupleTag tag, T value) { @Override public void outputWithTimestamp(TupleTag tag, T value, Instant timestamp) { outputReceiver.output( - tag, WindowedValues.of(value, timestamp, element.getWindows(), element.getPaneInfo())); + tag, + WindowedValues.of( + value, + timestamp, + element.getWindows(), + element.getPaneInfo(), + element.getRecordId(), + element.getRecordOffset(), + element.isDraining())); } @Override @@ -479,7 +498,13 @@ public void outputWindowedValue( outputReceiver.output( tag, WindowedValues.of( - value, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + value, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + element.isDraining())); } private void noteOutput() { diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFn.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFn.java index 0be0c6353783..146a60203f29 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFn.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFn.java @@ -59,6 +59,7 @@ public abstract class ProcessValueContext extends Context { /** Return the actual value being processed. */ public abstract InputT value(); + public abstract Boolean draining(); /** Return the timestamp associated with the value. */ public abstract Instant timestamp(); } @@ -75,6 +76,7 @@ public abstract class OnTriggerContext extends Context { /** Returns the {@link PaneInfo} for the trigger firing being processed. */ public abstract PaneInfo paneInfo(); + public abstract Boolean draining(); /** Output the given value in the current window. */ public abstract void output(OutputT value); } diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java index 1ca927f16caa..ae175a5f0117 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java @@ -94,17 +94,26 @@ private StateAccessorImpl stateAccessor(W window, StateStyle style) { } public ReduceFn.Context base(W window, StateStyle style) { - return new ContextImpl(stateAccessor(window, style)); + return new ContextImpl(stateAccessor(window, style), null); + } + + public ReduceFn.Context base( + W window, StateStyle style, Boolean draining) { + return new ContextImpl(stateAccessor(window, style), draining); } public ReduceFn.ProcessValueContext forValue( - W window, InputT value, Instant timestamp, StateStyle style) { - return new ProcessValueContextImpl(stateAccessor(window, style), value, timestamp); + W window, InputT value, Instant timestamp, StateStyle style, Boolean draining) { + return new ProcessValueContextImpl(stateAccessor(window, style), value, timestamp, draining); } public ReduceFn.OnTriggerContext forTrigger( - W window, PaneInfo paneInfo, StateStyle style, OnTriggerCallbacks callbacks) { - return new OnTriggerContextImpl(stateAccessor(window, style), paneInfo, callbacks); + W window, + PaneInfo paneInfo, + StateStyle style, + OnTriggerCallbacks callbacks, + Boolean draining) { + return new OnTriggerContextImpl(stateAccessor(window, style), paneInfo, callbacks, draining); } public ReduceFn.OnMergeContext forMerge( @@ -135,16 +144,19 @@ public TimersImpl(StateNamespace namespace) { @Override public void setTimer(Instant timestamp, TimeDomain timeDomain) { + // todo radoslaws@ should we pass draining bit to timerdata timerInternals.setTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain)); } @Override public void setTimer(Instant timestamp, Instant outputTimestamp, TimeDomain timeDomain) { + // todo radoslaws@ should we pass draining bit to timerdata timerInternals.setTimer(TimerData.of(namespace, timestamp, outputTimestamp, timeDomain)); } @Override public void deleteTimer(Instant timestamp, TimeDomain timeDomain) { + // todo radoslaws@ should we pass draining bit to timerdata timerInternals.deleteTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain)); } @@ -315,11 +327,13 @@ public Map accessInEachMergingWindow( private class ContextImpl extends ReduceFn.Context { private final StateAccessorImpl state; private final TimersImpl timers; + private final Boolean draining; - private ContextImpl(StateAccessorImpl state) { + private ContextImpl(StateAccessorImpl state, @Nullable Boolean draining) { reduceFn.super(); this.state = state; this.timers = new TimersImpl(state.namespace()); + this.draining = draining; } @Override @@ -354,14 +368,16 @@ private class ProcessValueContextImpl private final Instant timestamp; private final StateAccessorImpl state; private final TimersImpl timers; + private final Boolean draining; private ProcessValueContextImpl( - StateAccessorImpl state, InputT value, Instant timestamp) { + StateAccessorImpl state, InputT value, Instant timestamp, Boolean draining) { reduceFn.super(); this.state = state; this.value = value; this.timestamp = timestamp; this.timers = new TimersImpl(state.namespace()); + this.draining = draining; } @Override @@ -369,6 +385,11 @@ public K key() { return key; } + @Override + public Boolean draining() { + return draining; + } + @Override public W window() { return state.window(); @@ -405,14 +426,19 @@ private class OnTriggerContextImpl extends ReduceFn.OnTri private final PaneInfo paneInfo; private final OnTriggerCallbacks callbacks; private final TimersImpl timers; + private final Boolean draining; private OnTriggerContextImpl( - StateAccessorImpl state, PaneInfo paneInfo, OnTriggerCallbacks callbacks) { + StateAccessorImpl state, + PaneInfo paneInfo, + OnTriggerCallbacks callbacks, + Boolean draining) { reduceFn.super(); this.state = state; this.paneInfo = paneInfo; this.callbacks = callbacks; this.timers = new TimersImpl(state.namespace()); + this.draining = draining; } @Override @@ -420,6 +446,11 @@ public K key() { return key; } + @Override + public Boolean draining() { + return draining; + } + @Override public W window() { return state.window(); diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java index b08bd42b0b22..528a62f46c0a 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java @@ -374,7 +374,8 @@ public void processElements(Iterable> values) throws Excep for (W window : windowsToFire) { emit( contextFactory.base(window, StateStyle.DIRECT), - contextFactory.base(window, StateStyle.RENAMED)); + contextFactory.base(window, StateStyle.RENAMED), + null); } // We're all done with merging and emitting elements so can compress the activeWindow state. @@ -583,7 +584,11 @@ private void processElement(Map windowToMergeResult, WindowedValue for (W window : windows) { ReduceFn.ProcessValueContext directContext = contextFactory.forValue( - window, value.getValue(), value.getTimestamp(), StateStyle.DIRECT); + window, + value.getValue(), + value.getTimestamp(), + StateStyle.DIRECT, + value.isDraining()); if (triggerRunner.isClosed(directContext.state())) { // This window has already been closed. droppedDueToClosedWindow.inc(); @@ -601,7 +606,11 @@ private void processElement(Map windowToMergeResult, WindowedValue activeWindows.ensureWindowIsActive(window); ReduceFn.ProcessValueContext renamedContext = contextFactory.forValue( - window, value.getValue(), value.getTimestamp(), StateStyle.RENAMED); + window, + value.getValue(), + value.getTimestamp(), + StateStyle.RENAMED, + value.isDraining()); nonEmptyPanes.recordContent(renamedContext.state()); scheduleGarbageCollectionTimer(directContext); @@ -639,12 +648,15 @@ private class WindowActivation { // garbage collect the window. We'll consider any timer at or after the // end-of-window time to be a signal to garbage collect. public final boolean isGarbageCollection; + public final Boolean draining; WindowActivation( ReduceFn.Context directContext, - ReduceFn.Context renamedContext) { + ReduceFn.Context renamedContext, + Boolean draining) { this.directContext = directContext; this.renamedContext = renamedContext; + this.draining = draining; W window = directContext.window(); // The output watermark is before the end of the window if it is either unknown @@ -701,12 +713,13 @@ public void onTimers(Iterable timers) throws Exception { WindowTracing.debug( "ReduceFnRunner: Received timer key:{}; window:{}; data:{} with " - + "inputWatermark:{}; outputWatermark:{}", + + "inputWatermark:{}; outputWatermark:{}; draining:{}", key, window, timer, timerInternals.currentInputWatermarkTime(), - timerInternals.currentOutputWatermarkTime()); + timerInternals.currentOutputWatermarkTime(), + timer.getDraining()); // Processing time timers for an expired window are ignored, just like elements // that show up too late. Window GC is management by an event time timer @@ -722,11 +735,13 @@ public void onTimers(Iterable timers) throws Exception { continue; } + // pass draining bit ReduceFn.Context directContext = contextFactory.base(window, StateStyle.DIRECT); ReduceFn.Context renamedContext = contextFactory.base(window, StateStyle.RENAMED); - WindowActivation windowActivation = new WindowActivation(directContext, renamedContext); + WindowActivation windowActivation = + new WindowActivation(directContext, renamedContext, timer.getDraining()); windowActivations.put(window, windowActivation); // Perform prefetching of state to determine if the trigger should fire. @@ -757,11 +772,12 @@ public void onTimers(Iterable timers) throws Exception { if (windowActivation.isGarbageCollection) { WindowTracing.debug( - "ReduceFnRunner: Cleaning up for key:{}; window:{} with inputWatermark:{}; outputWatermark:{}", + "ReduceFnRunner: Cleaning up for key:{}; window:{} with inputWatermark:{}; outputWatermark:{}; draining:{}", key, directContext.window(), timerInternals.currentInputWatermarkTime(), - timerInternals.currentOutputWatermarkTime()); + timerInternals.currentOutputWatermarkTime(), + windowActivation.draining); boolean windowIsActiveAndOpen = windowActivation.windowIsActiveAndOpen(); if (windowIsActiveAndOpen) { @@ -774,7 +790,8 @@ public void onTimers(Iterable timers) throws Exception { directContext, renamedContext, true /* isFinished */, - windowActivation.isEndOfWindow); + windowActivation.isEndOfWindow, + windowActivation.draining); checkState(newHold == null, "Hold placed at %s despite isFinished being true.", newHold); } @@ -792,7 +809,7 @@ public void onTimers(Iterable timers) throws Exception { if (windowActivation.windowIsActiveAndOpen() && triggerRunner.shouldFire( directContext.window(), directContext.timers(), directContext.state())) { - emit(directContext, renamedContext); + emit(directContext, renamedContext, windowActivation.draining); } if (windowActivation.isEndOfWindow) { @@ -915,7 +932,8 @@ private void prefetchEmit( /** Emit if a trigger is ready to fire or timers require it, and cleanup state. */ private void emit( ReduceFn.Context directContext, - ReduceFn.Context renamedContext) + ReduceFn.Context renamedContext, + Boolean draining) throws Exception { checkState( triggerRunner.shouldFire( @@ -931,7 +949,7 @@ private void emit( // Run onTrigger to produce the actual pane contents. // As a side effect it will clear all element holds, but not necessarily any // end-of-window or garbage collection holds. - onTrigger(directContext, renamedContext, isFinished, false /*isEndOfWindow*/); + onTrigger(directContext, renamedContext, isFinished, false /*isEndOfWindow*/, draining); // Now that we've triggered, the pane is empty. nonEmptyPanes.clearPane(renamedContext.state()); @@ -989,7 +1007,8 @@ private void prefetchOnTrigger( final ReduceFn.Context directContext, ReduceFn.Context renamedContext, final boolean isFinished, - boolean isEndOfWindow) + boolean isEndOfWindow, + Boolean draining) throws Exception { // Extract the window hold, and as a side effect clear it. final WatermarkHold.OldAndNewHolds pair = @@ -1061,10 +1080,12 @@ private void prefetchOnTrigger( .setValue(KV.of(key, toOutput)) .setTimestamp(outputTimestamp) .setWindows(windows) + .setDraining(draining) .setPaneInfo(paneInfo) .setReceiver(outputter) .output(); - }); + }, + draining); reduceFn.onTrigger(renamedTriggerContext); } diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java index 3af90ea9a0a1..f54687bfa629 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java @@ -201,11 +201,13 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { Preconditions.checkNotNull(outputTimestamp, "outputTimestamp"); OnTimerArgumentProvider argumentProvider = - new OnTimerArgumentProvider<>(timerId, key, window, timestamp, outputTimestamp, timeDomain); + new OnTimerArgumentProvider<>( + timerId, key, window, timestamp, outputTimestamp, timeDomain, draining); invoker.invokeOnTimer(timerId, timerFamilyId, argumentProvider); } @@ -363,7 +365,8 @@ public void output( Collections.singletonList(window), PaneInfo.NO_FIRING, currentRecordId, - currentRecordOffset)); + currentRecordOffset, + null)); } } @@ -429,6 +432,11 @@ public InputT element() { return elem.getValue(); } + @Override + public Boolean draining() { + return elem.isDraining(); + } + @Override public T sideInput(PCollectionView view) { checkNotNull(view, "View passed to sideInput cannot be null"); @@ -524,7 +532,13 @@ public void outputWindowedValue( SimpleDoFnRunner.this.outputWindowedValue( tag, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + elem.isDraining())); } @Override @@ -725,6 +739,7 @@ private class OnTimerArgumentProvider extends DoFn.OnTime private final TimeDomain timeDomain; private final String timerId; private final KeyT key; + private final Boolean draining; private final OutputBuilderSupplier builderSupplier; /** Lazily initialized; should only be accessed via {@link #getNamespace()}. */ @@ -750,7 +765,8 @@ private OnTimerArgumentProvider( BoundedWindow window, Instant fireTimestamp, Instant timestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + Boolean draining) { fn.super(); this.timerId = timerId; this.window = window; @@ -758,13 +774,15 @@ private OnTimerArgumentProvider( this.timestamp = timestamp; this.timeDomain = timeDomain; this.key = key; + this.draining = draining; this.builderSupplier = OutputBuilderSuppliers.supplierForElement( WindowedValues.builder() .setValue(null) .setTimestamp(timestamp) .setWindow(window) - .setPaneInfo(PaneInfo.NO_FIRING)); + .setPaneInfo(PaneInfo.NO_FIRING) + .setDraining(draining)); } @Override @@ -772,6 +790,11 @@ public Instant timestamp() { return timestamp; } + @Override + public Boolean draining() { + return draining; + } + @Override public Instant fireTimestamp() { return fireTimestamp; @@ -1026,7 +1049,7 @@ public void outputWindowedValue( SimpleDoFnRunner.this.outputWindowedValue( tag, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset, null)); } @Override @@ -1304,7 +1327,7 @@ public void outputWindowedValue( SimpleDoFnRunner.this.outputWindowedValue( tag, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset, null)); } @Override diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java index 2be8071f9835..fec32485b827 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java @@ -116,7 +116,8 @@ public void onTimer( Instant timestamp, Instant outputTimestamp, TimeDomain timeDomain) { - underlying.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + underlying.onTimer( + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, null); } @Override diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java index 6af54da0a08b..7ddb146e0eb1 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java @@ -598,6 +598,7 @@ public String getErrorContext() { timerInternals.currentProcessingTime().plus(result.getContinuation().resumeDelay()); holdState.add(futureOutputWatermark); // Set a timer to continue processing this element. + // todo radoslws@ decide if draining should be set on timer timerInternals.setTimer( TimerInternals.TimerData.of( stateNamespace, wakeupTime, wakeupTime, TimeDomain.PROCESSING_TIME)); diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java index 52bbb3a306a0..e54f45be6bef 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java @@ -42,6 +42,7 @@ 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.base.MoreObjects; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; import org.joda.time.Instant; @@ -208,7 +209,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { if (timerId.equals(SORT_FLUSH_TIMER)) { onSortFlushTimer(window, stepContext.timerInternals().currentInputWatermarkTime()); @@ -232,7 +234,7 @@ public void onTimer( stepContext.timerInternals().currentInputWatermarkTime()); } else { doFnRunner.onTimer( - timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } } } diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java index 254e6f5fcf5b..fabc755aae71 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java @@ -23,9 +23,11 @@ import java.io.OutputStream; import java.util.Arrays; import java.util.List; +import org.apache.beam.sdk.coders.BooleanCoder; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderException; import org.apache.beam.sdk.coders.InstantCoder; +import org.apache.beam.sdk.coders.NullableCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.coders.StructuredCoder; import org.apache.beam.sdk.state.TimeDomain; @@ -188,6 +190,8 @@ abstract class TimerData implements Comparable { public abstract boolean getDeleted(); + public abstract @javax.annotation.Nullable Boolean getDraining(); + // When adding a new field, make sure to add it to the compareTo() method. /** Construct a {@link TimerData} for the given parameters. */ @@ -198,7 +202,7 @@ public static TimerData of( Instant outputTimestamp, TimeDomain domain) { return new AutoValue_TimerInternals_TimerData( - timerId, "", namespace, timestamp, outputTimestamp, domain, false); + timerId, "", namespace, timestamp, outputTimestamp, domain, false, null); } /** @@ -213,7 +217,7 @@ public static TimerData of( Instant outputTimestamp, TimeDomain domain) { return new AutoValue_TimerInternals_TimerData( - timerId, timerFamilyId, namespace, timestamp, outputTimestamp, domain, false); + timerId, timerFamilyId, namespace, timestamp, outputTimestamp, domain, false, null); } /** @@ -234,7 +238,20 @@ public TimerData deleted() { getTimestamp(), getOutputTimestamp(), getDomain(), - true); + true, + getDraining()); + } + + public TimerData draining(@Nullable Boolean draining) { + return new AutoValue_TimerInternals_TimerData( + getTimerId(), + getTimerFamilyId(), + getNamespace(), + getTimestamp(), + getOutputTimestamp(), + getDomain(), + getDeleted(), + draining); } /** @@ -272,7 +289,9 @@ public String stringKey() { + "/" + getTimerFamilyId() + ":" - + getTimerId(); + + getTimerId() + + ":" + + getDraining(); } } @@ -280,6 +299,7 @@ public String stringKey() { class TimerDataCoderV2 extends StructuredCoder { private static final StringUtf8Coder STRING_CODER = StringUtf8Coder.of(); private static final InstantCoder INSTANT_CODER = InstantCoder.of(); + private static final NullableCoder BOOLEAN_CODER = NullableCoder.of(BooleanCoder.of()); private final Coder windowCoder; public static TimerDataCoderV2 of(Coder windowCoder) { @@ -298,6 +318,7 @@ public void encode(TimerData timer, OutputStream outStream) throws CoderExceptio INSTANT_CODER.encode(timer.getTimestamp(), outStream); INSTANT_CODER.encode(timer.getOutputTimestamp(), outStream); STRING_CODER.encode(timer.getDomain().name(), outStream); + // we don't need to serialize draining bit } @Override @@ -309,6 +330,7 @@ public TimerData decode(InputStream inStream) throws CoderException, IOException Instant timestamp = INSTANT_CODER.decode(inStream); Instant outputTimestamp = INSTANT_CODER.decode(inStream); TimeDomain domain = TimeDomain.valueOf(STRING_CODER.decode(inStream)); + // we don't need to serialize draining bit return TimerData.of(timerId, timerFamilyId, namespace, timestamp, outputTimestamp, domain); } diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/WindowMatchers.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/WindowMatchers.java index 463cb9320237..fc945791a788 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/WindowMatchers.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/WindowMatchers.java @@ -59,7 +59,7 @@ public static Matcher> isWindowedValue( Matcher> windowsMatcher, Matcher paneInfoMatcher) { return new WindowedValueMatcher<>( - valueMatcher, timestampMatcher, windowsMatcher, paneInfoMatcher); + valueMatcher, timestampMatcher, windowsMatcher, paneInfoMatcher, Matchers.anything()); } public static Matcher> isWindowedValue( @@ -67,19 +67,27 @@ public static Matcher> isWindowedValue( Matcher timestampMatcher, Matcher> windowsMatcher) { return new WindowedValueMatcher<>( - valueMatcher, timestampMatcher, windowsMatcher, Matchers.anything()); + valueMatcher, timestampMatcher, windowsMatcher, Matchers.anything(), Matchers.anything()); } public static Matcher> isWindowedValue( Matcher valueMatcher, Matcher timestampMatcher) { return new WindowedValueMatcher<>( - valueMatcher, timestampMatcher, Matchers.anything(), Matchers.anything()); + valueMatcher, + timestampMatcher, + Matchers.anything(), + Matchers.anything(), + Matchers.anything()); } public static Matcher> isWindowedValue( Matcher valueMatcher) { return new WindowedValueMatcher<>( - valueMatcher, Matchers.anything(), Matchers.anything(), Matchers.anything()); + valueMatcher, + Matchers.anything(), + Matchers.anything(), + Matchers.anything(), + Matchers.anything()); } public static Matcher> isSingleWindowedValue( @@ -143,7 +151,25 @@ public static Matcher> isSingleWindowedValue( Matcher timestampMatcher, Matcher windowMatcher) { return new WindowedValueMatcher<>( - valueMatcher, timestampMatcher, Matchers.contains(windowMatcher), Matchers.anything()); + valueMatcher, + timestampMatcher, + Matchers.contains(windowMatcher), + Matchers.anything(), + Matchers.anything()); + } + + public static Matcher> isSingleWindowedValue( + Matcher valueMatcher, + Matcher timestampMatcher, + Matcher windowMatcher, + Matcher paneInfoMatcher, + Matcher drainMatcher) { + return new WindowedValueMatcher<>( + valueMatcher, + timestampMatcher, + Matchers.contains(windowMatcher), + paneInfoMatcher, + drainMatcher); } public static Matcher> isSingleWindowedValue( @@ -152,7 +178,11 @@ public static Matcher> isSingleWindowedValue( Matcher windowMatcher, Matcher paneInfoMatcher) { return new WindowedValueMatcher<>( - valueMatcher, timestampMatcher, Matchers.contains(windowMatcher), paneInfoMatcher); + valueMatcher, + timestampMatcher, + Matchers.contains(windowMatcher), + paneInfoMatcher, + Matchers.anything()); } public static Matcher intervalWindow(long start, long end) { @@ -203,16 +233,19 @@ private static class WindowedValueMatcher extends TypeSafeMatcher timestampMatcher; private Matcher> windowsMatcher; private Matcher paneInfoMatcher; + private Matcher drainMatcher; private WindowedValueMatcher( Matcher valueMatcher, Matcher timestampMatcher, Matcher> windowsMatcher, - Matcher paneInfoMatcher) { + Matcher paneInfoMatcher, + Matcher drainMatcher) { this.valueMatcher = valueMatcher; this.timestampMatcher = timestampMatcher; this.windowsMatcher = windowsMatcher; this.paneInfoMatcher = paneInfoMatcher; + this.drainMatcher = drainMatcher; } @Override @@ -226,6 +259,8 @@ public void describeTo(Description description) { .appendValue(windowsMatcher) .appendText(", ") .appendValue(paneInfoMatcher) + .appendText(", ") + .appendValue(drainMatcher) .appendText(")"); } @@ -234,7 +269,8 @@ protected boolean matchesSafely(WindowedValue windowedValue) { return valueMatcher.matches(windowedValue.getValue()) && timestampMatcher.matches(windowedValue.getTimestamp()) && windowsMatcher.matches(windowedValue.getWindows()) - && paneInfoMatcher.matches(windowedValue.getPaneInfo()); + && paneInfoMatcher.matches(windowedValue.getPaneInfo()) + && drainMatcher.matches(windowedValue.isDraining()); } } } diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java index 55791d582717..51ba557749e2 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java @@ -89,7 +89,6 @@ }) public class ReduceFnTester { private static final String KEY = "TEST_KEY"; - private final TestInMemoryStateInternals stateInternals = new TestInMemoryStateInternals<>(KEY); private final InMemoryTimerInternals timerInternals = new InMemoryTimerInternals(); diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java index 40bcb70e32c6..52e241b0e92c 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java @@ -140,7 +140,8 @@ public void testOnTimerExceptionsWrappedAsUserCodeException() { GlobalWindow.INSTANCE, new Instant(0), new Instant(0), - TimeDomain.EVENT_TIME); + TimeDomain.EVENT_TIME, + null); } /** @@ -265,7 +266,8 @@ public void testOnTimerCalled() { GlobalWindow.INSTANCE, currentTime.plus(offset), currentTime.plus(offset), - TimeDomain.EVENT_TIME); + TimeDomain.EVENT_TIME, + null); assertThat( fn.onTimerInvocations, @@ -592,7 +594,8 @@ public void testOnTimerAllowedSkew() { GlobalWindow.INSTANCE, new Instant(0), new Instant(0), - TimeDomain.EVENT_TIME); + TimeDomain.EVENT_TIME, + null); } @Test @@ -624,7 +627,8 @@ public void testOnTimerNoSkew() { GlobalWindow.INSTANCE, new Instant(0), new Instant(0), - TimeDomain.EVENT_TIME); + TimeDomain.EVENT_TIME, + null); }); assertThat(exception.getCause(), isA(IllegalArgumentException.class)); diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java index b7dda599b8d2..5d0ba9d77c99 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java @@ -64,6 +64,7 @@ import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; 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.Duration; import org.joda.time.Instant; import org.junit.Before; @@ -351,15 +352,17 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { firedTimers.add( TimerData.of( - timerId, - timerFamilyId, - StateNamespaces.window(IntervalWindow.getCoder(), (IntervalWindow) window), - timestamp, - outputTimestamp, - timeDomain)); + timerId, + timerFamilyId, + StateNamespaces.window(IntervalWindow.getCoder(), (IntervalWindow) window), + timestamp, + outputTimestamp, + timeDomain) + .draining(draining)); } @Override diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/StatefulDoFnRunnerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/StatefulDoFnRunnerTest.java index 4ed4f09fe62b..813d893c826b 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/StatefulDoFnRunnerTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/StatefulDoFnRunnerTest.java @@ -468,7 +468,8 @@ private static void advanceInputWatermark( window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain()); + timer.getDomain(), + timer.getDraining()); } } diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java index f1ec36564a4b..cc0cf4e1a0b0 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java @@ -26,6 +26,7 @@ import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -73,10 +74,12 @@ public void onTimer( final BoundedWindow window, final Instant timestamp, final Instant outputTimestamp, - final TimeDomain timeDomain) { + final TimeDomain timeDomain, + @Nullable Boolean draining) { try (Closeable ignored = MetricsEnvironment.scopedMetricsContainer(container.getMetricsContainer(stepName))) { - delegate.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + delegate.onTimer( + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } catch (IOException e) { throw new RuntimeException(e); } diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkStatefulDoFnFunction.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkStatefulDoFnFunction.java index 3aa5a0802b38..a387633f6ccd 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkStatefulDoFnFunction.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkStatefulDoFnFunction.java @@ -250,7 +250,8 @@ private void fireTimer( window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain()); + timer.getDomain(), + null); } @Override diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java index d5a7ff035efc..8b1b2d08cd49 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java @@ -1001,7 +1001,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { Object timerKey = keyForTimer.get(); Preconditions.checkNotNull(timerKey, "Key for timer needs to be set before calling onTimer"); Preconditions.checkNotNull(remoteBundle, "Call to onTimer outside of a bundle"); @@ -1034,7 +1035,8 @@ public void onTimer( timestamp, outputTimestamp, // TODO: Support propagating the PaneInfo through. - PaneInfo.NO_FIRING); + PaneInfo.NO_FIRING, + draining); try { timerReceiver.accept(timerValue); } catch (Exception e) { diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java index 0c2ba87c4baa..5cce6f9fab3b 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java @@ -104,7 +104,7 @@ static final class Timer implements BufferedElement { @Override public void processWith(DoFnRunner doFnRunner) { doFnRunner.onTimer( - timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, null); } @Override diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java index 10a3182f90f4..2025e6c4e402 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java @@ -234,7 +234,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { minBufferedElementTimestamp = Math.min(outputTimestamp.getMillis(), minBufferedElementTimestamp); diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java index c2556d7229ba..86e99bbb8017 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java @@ -507,7 +507,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { if ("cleanup".equals(timerId)) { holdState.clear(); diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperatorTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperatorTest.java index 982871e59f14..8447fae34178 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperatorTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperatorTest.java @@ -534,7 +534,8 @@ public void testWatermarkHandling() throws Exception { windowedValue.getWindows(), timestamp, timestamp, - PaneInfo.NO_FIRING), + PaneInfo.NO_FIRING, + null), TimerInternals.TimerData.of( "", TimerReceiverFactory.encodeToTimerDataTimerId("transform", timerId), @@ -875,7 +876,8 @@ private void testEnsureDeferredStateCleanupTimerFiring(boolean withCheckpointing windowedValue.getWindows(), window.maxTimestamp(), window.maxTimestamp(), - PaneInfo.NO_FIRING), + PaneInfo.NO_FIRING, + null), userTimer2); assertThat(testHarness.numEventTimeTimers(), is(1)); diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java index 10b41bb5b5ba..7026396a42b9 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java @@ -1378,6 +1378,11 @@ public T getValue() { return value; } + @Override + public @Nullable Boolean isDraining() { + return null; + } + @Override public Instant getTimestamp() { return BoundedWindow.TIMESTAMP_MIN_VALUE; diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/RedistributeByKeyOverrideFactory.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/RedistributeByKeyOverrideFactory.java index 4375cc5adcfe..47ff5b764910 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/RedistributeByKeyOverrideFactory.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/RedistributeByKeyOverrideFactory.java @@ -135,6 +135,7 @@ public Duration getAllowedTimestampSkew() { public void processElement( @Element KV> kv, OutputReceiver> outputReceiver) { + // todo #33176 specify additional metadata in the future outputReceiver .builder(KV.of(kv.getKey(), kv.getValue().getValue())) .setTimestamp(kv.getValue().getTimestamp()) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowProcessFnRunner.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowProcessFnRunner.java index 741716b82848..59250e4f15d3 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowProcessFnRunner.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowProcessFnRunner.java @@ -34,6 +34,7 @@ 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.Iterables; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -117,7 +118,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { throw new UnsupportedOperationException("Unsupported for ProcessFn"); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowFnRunner.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowFnRunner.java index 1909a73dc8b9..b30525ee068a 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowFnRunner.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowFnRunner.java @@ -26,6 +26,7 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.util.WindowedValueReceiver; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -81,7 +82,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { throw new UnsupportedOperationException( String.format("Timers are not supported by %s", GroupAlsoByWindowFn.class.getSimpleName())); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java index 00c8192b1e40..d3f280085155 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java @@ -374,7 +374,8 @@ private void processUserTimer(TimerData timer) throws Exception { window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain()); + timer.getDomain(), + timer.getDraining()); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java index 2a4b111af225..31729c3b71a9 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java @@ -110,6 +110,7 @@ import org.apache.beam.sdk.io.gcp.bigquery.BigQuerySinkMetrics; import org.apache.beam.sdk.metrics.MetricsEnvironment; import org.apache.beam.sdk.util.construction.CoderTranslation; +import org.apache.beam.sdk.values.WindowedValues; 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.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheStats; @@ -163,6 +164,7 @@ public final class StreamingDataflowWorker { private static final Random CLIENT_ID_GENERATOR = new Random(); private static final String CHANNELZ_PATH = "/channelz"; private static final String BEAM_FN_API_EXPERIMENT = "beam_fn_api"; + private static final String ELEMENT_METADATA_SUPPORTED_EXPERIMENT = "element_metadata_supported"; private static final String STREAMING_ENGINE_USE_JOB_SETTINGS_FOR_HEARTBEAT_POOL_EXPERIMENT = "streaming_engine_use_job_settings_for_heartbeat_pool"; // Experiment make the monitor within BoundedQueueExecutor fair @@ -261,6 +263,7 @@ private StreamingDataflowWorker( serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies) -> computationStateCache .get(processingContext.computationId()) @@ -273,6 +276,7 @@ private StreamingDataflowWorker( serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies); }), ChannelCachingRemoteStubFactory.create(options.getGcpCredential(), channelCache), @@ -815,6 +819,9 @@ public static void main(String[] args) throws Exception { validateWorkerOptions(options); CoderTranslation.verifyModelCodersRegistered(); + if (DataflowRunner.hasExperiment(options, ELEMENT_METADATA_SUPPORTED_EXPERIMENT)) { + WindowedValues.FullWindowedValueCoder.setMetadataSupported(); + } LOG.debug("Creating StreamingDataflowWorker from options: {}", options); StreamingDataflowWorker worker = StreamingDataflowWorker.fromOptions(options); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunner.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunner.java index e5d8a18be76f..b35f0e7da3cb 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunner.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunner.java @@ -38,6 +38,7 @@ 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.collect.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -142,7 +143,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { throw new UnsupportedOperationException( "Attempt to deliver a timer to a DoFn, but timers are not supported in Dataflow."); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index b24ca561495c..82d18b0ab788 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -195,6 +195,10 @@ public boolean workIsFailed() { return work != null && work.isFailed(); } + public @javax.annotation.Nullable Boolean getDrainMode() { + return work != null ? work.getDrainMode() : null; + } + public boolean offsetBasedDeduplicationSupported() { return activeReader != null && activeReader.getCurrentSource().offsetBasedDeduplicationSupported(); @@ -816,7 +820,10 @@ public TimerData getNextFiredTimer(Coder windowCode .transform( timer -> WindmillTimerInternals.windmillTimerToTimerData( - WindmillNamespacePrefix.SYSTEM_NAMESPACE_PREFIX, timer, windowCoder)) + WindmillNamespacePrefix.SYSTEM_NAMESPACE_PREFIX, + timer, + windowCoder, + getDrainMode())) .iterator(); } @@ -876,7 +883,10 @@ public TimerData getNextFiredUserTimer(Coder window .transform( timer -> WindmillTimerInternals.windmillTimerToTimerData( - WindmillNamespacePrefix.USER_NAMESPACE_PREFIX, timer, windowCoder)) + WindmillNamespacePrefix.USER_NAMESPACE_PREFIX, + timer, + windowCoder, + getDrainMode())) .iterator()); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java index b5b723adb2b5..6a556721cdd8 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java @@ -24,6 +24,7 @@ import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -84,7 +85,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { throw new UnsupportedOperationException( "Attempt to deliver a timer to a DoFn, but timers are not supported in Dataflow."); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java index e031d1bb50eb..ff931b1bf6be 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java @@ -24,6 +24,7 @@ import java.io.InputStream; import java.util.Collection; import java.util.Map; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.runners.dataflow.util.CloudObject; import org.apache.beam.runners.dataflow.worker.util.common.worker.NativeReader; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; @@ -117,6 +118,15 @@ protected WindowedValue decodeMessage(Windmill.Message message) throws IOExce Collection windows = WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata()); PaneInfo paneInfo = WindmillSink.decodeMetadataPane(message.getMetadata()); + Boolean drainingValueFromUpstream = null; + if (WindowedValues.WindowedValueCoder.isMetadataSupported()) { + BeamFnApi.Elements.ElementMetadata elementMetadata = + WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata()); + drainingValueFromUpstream = + elementMetadata.hasDrain() + ? (elementMetadata.getDrain() == BeamFnApi.Elements.DrainMode.Enum.DRAINING) + : null; + } if (valueCoder instanceof KvCoder) { KvCoder kvCoder = (KvCoder) valueCoder; InputStream key = context.getSerializedKey().newInput(); @@ -125,10 +135,20 @@ protected WindowedValue decodeMessage(Windmill.Message message) throws IOExce @SuppressWarnings("unchecked") T result = (T) KV.of(decode(kvCoder.getKeyCoder(), key), decode(kvCoder.getValueCoder(), data)); - return WindowedValues.of(result, timestampMillis, windows, paneInfo); + // todo #33176 propagate metadata to windowed value + return WindowedValues.of( + result, timestampMillis, windows, paneInfo, null, null, drainingValueFromUpstream); } else { notifyElementRead(data.available() + metadata.available()); - return WindowedValues.of(decode(valueCoder, data), timestampMillis, windows, paneInfo); + // todo #33176 propagate metadata to windowed value + return WindowedValues.of( + decode(valueCoder, data), + timestampMillis, + windows, + paneInfo, + null, + null, + drainingValueFromUpstream); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java index cee4894e3d68..4f67dfcef319 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java @@ -24,6 +24,7 @@ import java.util.Collection; import java.util.List; import java.util.Objects; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.runners.core.KeyedWorkItem; import org.apache.beam.runners.core.KeyedWorkItemCoder; import org.apache.beam.runners.core.TimerInternals.TimerData; @@ -60,6 +61,7 @@ public class WindmillKeyedWorkItem implements KeyedWorkItem private final Windmill.WorkItem workItem; private final K key; + private final @javax.annotation.Nullable Boolean drainMode; private final transient Coder windowCoder; private final transient Coder> windowsCoder; @@ -70,12 +72,14 @@ public WindmillKeyedWorkItem( Windmill.WorkItem workItem, Coder windowCoder, Coder> windowsCoder, - Coder valueCoder) { + Coder valueCoder, + @javax.annotation.Nullable Boolean drainMode) { this.key = key; this.workItem = workItem; this.windowCoder = windowCoder; this.windowsCoder = windowsCoder; this.valueCoder = valueCoder; + this.drainMode = drainMode; } @Override @@ -93,7 +97,10 @@ public Iterable timersIterable() { .transform( timer -> WindmillTimerInternals.windmillTimerToTimerData( - WindmillNamespacePrefix.SYSTEM_NAMESPACE_PREFIX, timer, windowCoder)); + WindmillNamespacePrefix.SYSTEM_NAMESPACE_PREFIX, + timer, + windowCoder, + drainMode)); } @Override @@ -108,10 +115,22 @@ public Iterable> elementsIterable() { Collection windows = WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata()); PaneInfo paneInfo = WindmillSink.decodeMetadataPane(message.getMetadata()); - + // Draining value is based on upstream data + Boolean drainingValueFromUpstream = null; + if (WindowedValues.WindowedValueCoder.isMetadataSupported()) { + BeamFnApi.Elements.ElementMetadata elementMetadata = + WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata()); + drainingValueFromUpstream = + elementMetadata.hasDrain() + ? (elementMetadata.getDrain() + == BeamFnApi.Elements.DrainMode.Enum.DRAINING) + : null; + } InputStream inputStream = message.getData().newInput(); ElemT value = valueCoder.decode(inputStream, Coder.Context.OUTER); - return WindowedValues.of(value, timestamp, windows, paneInfo); + // todo #33176 specify additional metadata in the future + return WindowedValues.of( + value, timestamp, windows, paneInfo, null, null, drainingValueFromUpstream); } catch (IOException e) { throw new RuntimeException(e); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java index 7cb6f2223472..6d3484f6642d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java @@ -26,9 +26,11 @@ import java.util.Collection; import java.util.HashMap; import java.util.Map; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.runners.dataflow.util.CloudObject; import org.apache.beam.runners.dataflow.worker.util.common.worker.Sink; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; +import org.apache.beam.sdk.coders.ByteArrayCoder; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.options.PipelineOptions; @@ -40,6 +42,7 @@ import org.apache.beam.sdk.values.ValueWithRecordId; import org.apache.beam.sdk.values.ValueWithRecordId.ValueWithRecordIdCoder; import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.sdk.values.WindowedValues.FullWindowedValueCoder; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; @@ -75,11 +78,20 @@ private static ByteString encodeMetadata( ByteStringOutputStream stream, Coder> windowsCoder, Collection windows, - PaneInfo paneInfo) + PaneInfo paneInfo, + BeamFnApi.Elements.ElementMetadata metadata) throws IOException { try { - PaneInfoCoder.INSTANCE.encode(paneInfo, stream); - windowsCoder.encode(windows, stream, Coder.Context.OUTER); + // element metadata is behind the experiment + boolean elementMetadata = WindowedValues.WindowedValueCoder.isMetadataSupported(); + if (elementMetadata) { + PaneInfoCoder.INSTANCE.encode(paneInfo.withElementMetadata(true), stream); + windowsCoder.encode(windows, stream); + ByteArrayCoder.of().encode(metadata.toByteArray(), stream, Coder.Context.OUTER); + } else { + PaneInfoCoder.INSTANCE.encode(paneInfo, stream); + windowsCoder.encode(windows, stream, Coder.Context.OUTER); + } return stream.toByteStringAndReset(); } catch (Exception e) { stream.reset(); @@ -90,10 +102,11 @@ private static ByteString encodeMetadata( public static ByteString encodeMetadata( Coder> windowsCoder, Collection windows, - PaneInfo paneInfo) + PaneInfo paneInfo, + BeamFnApi.Elements.ElementMetadata metadata) throws IOException { ByteStringOutputStream stream = new ByteStringOutputStream(); - return encodeMetadata(stream, windowsCoder, windows, paneInfo); + return encodeMetadata(stream, windowsCoder, windows, paneInfo, metadata); } public static PaneInfo decodeMetadataPane(ByteString metadata) throws IOException { @@ -101,12 +114,26 @@ public static PaneInfo decodeMetadataPane(ByteString metadata) throws IOExceptio return PaneInfoCoder.INSTANCE.decode(inStream); } + public static BeamFnApi.Elements.ElementMetadata decodeAdditionalMetadata( + Coder> windowsCoder, ByteString metadata) + throws IOException { + InputStream inStream = metadata.newInput(); + PaneInfo paneInfo = PaneInfoCoder.INSTANCE.decode(inStream); + windowsCoder.decode(inStream); + if (paneInfo.isElementMetadata()) { + return BeamFnApi.Elements.ElementMetadata.parseFrom(ByteArrayCoder.of().decode(inStream)); + } else { + // empty + return BeamFnApi.Elements.ElementMetadata.newBuilder().build(); + } + } + public static Collection decodeMetadataWindows( Coder> windowsCoder, ByteString metadata) throws IOException { InputStream inStream = metadata.newInput(); PaneInfoCoder.INSTANCE.decode(inStream); - return windowsCoder.decode(inStream, Coder.Context.OUTER); + return windowsCoder.decode(inStream); } /** A {@link SinkFactory.Registrar} for windmill sinks. */ @@ -184,8 +211,12 @@ private ByteString encode(Coder coder, EncodeT object) throws public long add(WindowedValue data) throws IOException { ByteString key, value; ByteString id = ByteString.EMPTY; + // todo #33176 specify additional metadata in the future + BeamFnApi.Elements.ElementMetadata additionalMetadata = + BeamFnApi.Elements.ElementMetadata.newBuilder().build(); ByteString metadata = - encodeMetadata(stream, windowsCoder, data.getWindows(), data.getPaneInfo()); + encodeMetadata( + stream, windowsCoder, data.getWindows(), data.getPaneInfo(), additionalMetadata); if (valueCoder instanceof KvCoder) { KvCoder kvCoder = (KvCoder) valueCoder; KV kv = (KV) data.getValue(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java index 1dbc7b005345..a739ca7fcb1a 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java @@ -298,7 +298,10 @@ static Timer timerDataToWindmillTimer( } public static TimerData windmillTimerToTimerData( - WindmillNamespacePrefix prefix, Timer timer, Coder windowCoder) { + WindmillNamespacePrefix prefix, + Timer timer, + Coder windowCoder, + @javax.annotation.Nullable Boolean draining) { // The tag is a path-structure string but cheaper to parse than a proper URI. It follows // this pattern, where no component but the ID can contain a slash @@ -386,12 +389,13 @@ public static TimerData windmillTimerToTimerData( StateNamespace namespace = StateNamespaces.fromString(namespaceString, windowCoder); return TimerData.of( - timerId, - timerFamily, - namespace, - timestamp, - outputTimestamp, - timerTypeToTimeDomain(timer.getType())); + timerId, + timerFamily, + namespace, + timestamp, + outputTimestamp, + timerTypeToTimeDomain(timer.getType())) + .draining(draining); } private static boolean useNewTimerTagEncoding(TimerData timerData) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java index d91a5412b917..f4a6eec61cbf 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java @@ -119,7 +119,8 @@ public NativeReaderIterator>> iterator() throw final K key = keyCoder.decode(context.getSerializedKey().newInput(), Coder.Context.OUTER); final WorkItem workItem = context.getWorkItem(); KeyedWorkItem keyedWorkItem = - new WindmillKeyedWorkItem<>(key, workItem, windowCoder, windowsCoder, valueCoder); + new WindmillKeyedWorkItem<>( + key, workItem, windowCoder, windowsCoder, valueCoder, context.getDrainMode()); final boolean isEmptyWorkItem = (Iterables.isEmpty(keyedWorkItem.timersIterable()) && Iterables.isEmpty(keyedWorkItem.elementsIterable())); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java index 8b41a2d13219..8a1a864ed4ea 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java @@ -28,6 +28,7 @@ import java.util.Optional; import java.util.function.Consumer; import java.util.function.Supplier; +import javax.annotation.Nullable; import javax.annotation.concurrent.NotThreadSafe; import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.Pair; import org.apache.beam.runners.dataflow.worker.ActiveMessageMetadata; @@ -78,12 +79,14 @@ public final class Work implements RefreshableWork { private volatile TimedState currentState; private volatile boolean isFailed; private volatile String processingThreadName = ""; + private final @Nullable Boolean drainMode; private Work( WorkItem workItem, long serializedWorkItemSize, Watermarks watermarks, ProcessingContext processingContext, + @Nullable Boolean drainMode, Supplier clock) { this.shardedKey = ShardedKey.create(workItem.getKey(), workItem.getShardingKey()); this.workItem = workItem; @@ -91,6 +94,7 @@ private Work( this.processingContext = processingContext; this.watermarks = watermarks; this.clock = clock; + this.drainMode = drainMode; this.startTime = clock.get(); Preconditions.checkState(EMPTY_ENUM_MAP.isEmpty()); // Create by passing EMPTY_ENUM_MAP to avoid recreating @@ -110,8 +114,10 @@ public static Work create( long serializedWorkItemSize, Watermarks watermarks, ProcessingContext processingContext, + @Nullable Boolean drainMode, Supplier clock) { - return new Work(workItem, serializedWorkItemSize, watermarks, processingContext, clock); + return new Work( + workItem, serializedWorkItemSize, watermarks, processingContext, drainMode, clock); } public static ProcessingContext createProcessingContext( @@ -207,6 +213,10 @@ public State getState() { return currentState.state(); } + public @Nullable Boolean getDrainMode() { + return drainMode; + } + public void setState(State state) { Instant now = clock.get(); totalDurationPerState.compute( diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/SingleSourceWorkerHarness.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/SingleSourceWorkerHarness.java index 95023d117299..ce6c6a050861 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/SingleSourceWorkerHarness.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/SingleSourceWorkerHarness.java @@ -150,6 +150,7 @@ private void streamingEngineDispatchLoop( (computationId, inputDataWatermark, synchronizedProcessingTime, + drainMode, workItem, serializedWorkItemSize, getWorkStreamLatencies) -> @@ -173,6 +174,7 @@ private void streamingEngineDispatchLoop( getDataClient, workCommitter::commit, heartbeatSender), + drainMode, getWorkStreamLatencies); })); try { @@ -232,6 +234,7 @@ private void applianceDispatchLoop(Supplier getWorkFn) watermarks.setOutputDataWatermark(workItem.getOutputDataWatermark()).build(), Work.createProcessingContext( computationId, getDataClient, workCommitter::commit, heartbeatSender), + computationWork.getDrainMode(), /* getWorkStreamLatencies= */ ImmutableList.of()); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ValueInEmptyWindows.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ValueInEmptyWindows.java index a51c9ed419e1..90b31b974e2f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ValueInEmptyWindows.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ValueInEmptyWindows.java @@ -59,6 +59,11 @@ public PaneInfo getPaneInfo() { return null; } + @Override + public @Nullable Boolean isDraining() { + return null; + } + @Override public Iterable> explodeWindows() { return Collections.emptyList(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GetWorkResponseChunkAssembler.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GetWorkResponseChunkAssembler.java index 0ebb4726d3a1..017148e40e9f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GetWorkResponseChunkAssembler.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GetWorkResponseChunkAssembler.java @@ -124,7 +124,8 @@ private static ComputationMetadata fromProto( metadataProto.getComputationId(), WindmillTimeUtils.windmillToHarnessWatermark(metadataProto.getInputDataWatermark()), WindmillTimeUtils.windmillToHarnessWatermark( - metadataProto.getDependentRealtimeInputWatermark())); + metadataProto.getDependentRealtimeInputWatermark()), + metadataProto.getDrainMode()); } abstract String computationId(); @@ -132,6 +133,8 @@ private static ComputationMetadata fromProto( abstract @Nullable Instant inputDataWatermark(); abstract @Nullable Instant synchronizedProcessingTime(); + + abstract @Nullable Boolean drainMode(); } @AutoValue diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java index 2712bf1bd33d..8eb4c51a2b49 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java @@ -281,6 +281,7 @@ private void consumeAssembledWorkItem(AssembledWorkItem assembledWorkItem) { assembledWorkItem.bufferedSize(), createWatermarks(workItem, metadata), createProcessingContext(metadata.computationId()), + metadata.drainMode(), assembledWorkItem.latencyAttributions()); budgetTracker.recordBudgetReceived(assembledWorkItem.bufferedSize()); GetWorkBudget extension = budgetTracker.computeBudgetExtension(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetWorkStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetWorkStream.java index ae7ce85e13a8..58407ad8147f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetWorkStream.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetWorkStream.java @@ -203,6 +203,7 @@ private void consumeAssembledWorkItem(AssembledWorkItem assembledWorkItem) { assembledWorkItem.computationMetadata().computationId(), assembledWorkItem.computationMetadata().inputDataWatermark(), assembledWorkItem.computationMetadata().synchronizedProcessingTime(), + assembledWorkItem.computationMetadata().drainMode(), assembledWorkItem.workItem(), assembledWorkItem.bufferedSize(), assembledWorkItem.latencyAttributions()); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemReceiver.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemReceiver.java index e2f69585e48f..160bdaa07f89 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemReceiver.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemReceiver.java @@ -30,6 +30,7 @@ void receiveWork( String computation, @Nullable Instant inputDataWatermark, @Nullable Instant synchronizedProcessingTime, + @Nullable Boolean drainMode, Windmill.WorkItem workItem, long serializedWorkItemSize, ImmutableList getWorkStreamLatencies); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemScheduler.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemScheduler.java index b9d31fbe501d..4195dbca4bb5 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemScheduler.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemScheduler.java @@ -18,6 +18,7 @@ package org.apache.beam.runners.dataflow.worker.windmill.work; import javax.annotation.CheckReturnValue; +import javax.annotation.Nullable; import org.apache.beam.runners.dataflow.worker.streaming.Watermarks; import org.apache.beam.runners.dataflow.worker.streaming.Work; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.LatencyAttribution; @@ -35,6 +36,7 @@ public interface WorkItemScheduler { * @param workItem {@link WorkItem} to be processed. * @param watermarks processing watermarks for the workItem. * @param processingContext for processing the workItem. + * @param drainMode is job is draining. * @param getWorkStreamLatencies Latencies per processing stage for the WorkItem for reporting * back to Streaming Engine backend. */ @@ -43,5 +45,6 @@ void scheduleWork( long serializedWorkItemSize, Watermarks watermarks, Work.ProcessingContext processingContext, + @Nullable Boolean drainMode, ImmutableList getWorkStreamLatencies); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java index a4cd5d6d8a6b..0f0b16a09167 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java @@ -210,10 +210,12 @@ public void scheduleWork( long serializedWorkItemSize, Watermarks watermarks, Work.ProcessingContext processingContext, + @javax.annotation.Nullable Boolean drainMode, ImmutableList getWorkStreamLatencies) { computationState.activateWork( ExecutableWork.create( - Work.create(workItem, serializedWorkItemSize, watermarks, processingContext, clock), + Work.create( + workItem, serializedWorkItemSize, watermarks, processingContext, drainMode, clock), work -> processWork(computationState, work, getWorkStreamLatencies))); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java index dd13d5b55930..0daf96d9400b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java @@ -269,6 +269,7 @@ public boolean awaitTermination(int time, TimeUnit unit) throws InterruptedExcep computationWork.getComputationId(), inputDataWatermark, Instant.now(), + null, workItem, workItem.getSerializedSize(), ImmutableList.of( diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java index a60535dfbd69..c9aaef4d2e2c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java @@ -361,6 +361,7 @@ private static ExecutableWork createMockWork( Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build(), Work.createProcessingContext( computationId, new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, Instant::now), processWorkFn); } @@ -3541,6 +3542,7 @@ public void testLatencyAttributionProtobufsPopulated() { new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, clock); clock.sleep(Duration.millis(10)); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java index c89a031b3728..0cf2dc69111e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java @@ -19,16 +19,14 @@ import static org.apache.beam.sdk.TestUtils.KvMatcher.isKv; import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.containsInAnyOrder; -import static org.hamcrest.Matchers.emptyIterable; -import static org.hamcrest.Matchers.equalTo; -import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.*; import static org.mockito.Mockito.when; import java.io.IOException; import java.util.Arrays; import java.util.Collection; import java.util.List; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.runners.core.DoFnRunner; import org.apache.beam.runners.core.DoFnRunners; import org.apache.beam.runners.core.InMemoryStateInternals; @@ -176,7 +174,12 @@ private void addElement( valueCoder.encode(value, dataOutput, Context.OUTER); messageBundle .addMessagesBuilder() - .setMetadata(WindmillSink.encodeMetadata(windowsCoder, windows, PaneInfo.NO_FIRING)) + .setMetadata( + WindmillSink.encodeMetadata( + windowsCoder, + windows, + PaneInfo.NO_FIRING, + BeamFnApi.Elements.ElementMetadata.newBuilder().build())) .setData(dataOutput.toByteString()) .setTimestamp(WindmillTimeUtils.harnessToWindmillTimestamp(timestamp)); } @@ -188,7 +191,17 @@ private WindowedValue> createValue( return new ValueInEmptyWindows<>( (KeyedWorkItem) new WindmillKeyedWorkItem<>( - KEY, workItem.build(), windowCoder, wildcardWindowsCoder, valueCoder)); + KEY, workItem.build(), windowCoder, wildcardWindowsCoder, valueCoder, false)); + } + + private WindowedValue> createDrainingValue( + WorkItem.Builder workItem, Coder valueCoder) { + @SuppressWarnings({"unchecked", "rawtypes"}) + Coder> wildcardWindowsCoder = (Coder) windowsCoder; + return new ValueInEmptyWindows<>( + (KeyedWorkItem) + new WindmillKeyedWorkItem<>( + KEY, workItem.build(), windowCoder, wildcardWindowsCoder, valueCoder, true)); } @Test @@ -249,6 +262,78 @@ public void testFixedWindows() throws Exception { equalTo(window(10, 20))))); } + @Test + public void testFixedWindowsWithDraining() throws Exception { + TupleTag>> outputTag = new TupleTag<>(); + ListOutputManager outputManager = new ListOutputManager(); + DoFnRunner, KV>> runner = + makeRunner( + output -> outputManager.output(outputTag, output), + WindowingStrategy.of(FixedWindows.of(Duration.millis(10)))); + + when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(0)); + + runner.startBundle(); + + WorkItem.Builder workItem1 = WorkItem.newBuilder(); + workItem1.setKey(ByteString.copyFromUtf8(KEY)); + workItem1.setWorkToken(WORK_TOKEN); + InputMessageBundle.Builder messageBundle = workItem1.addMessageBundlesBuilder(); + messageBundle.setSourceComputationId(SOURCE_COMPUTATION_ID); + + Coder valueCoder = StringUtf8Coder.of(); + addElement(messageBundle, Arrays.asList(window(0, 10)), new Instant(1), valueCoder, "v1"); + addElement(messageBundle, Arrays.asList(window(0, 10)), new Instant(2), valueCoder, "v2"); + addElement(messageBundle, Arrays.asList(window(0, 10)), new Instant(0), valueCoder, "v0"); + addElement(messageBundle, Arrays.asList(window(10, 20)), new Instant(13), valueCoder, "v3"); + runner.processElement(createValue(workItem1, valueCoder)); + + runner.finishBundle(); + runner.startBundle(); + + WorkItem.Builder workItem2 = WorkItem.newBuilder(); + workItem2.setKey(ByteString.copyFromUtf8(KEY)); + workItem2.setWorkToken(WORK_TOKEN); + InputMessageBundle.Builder messageBundle2 = workItem1.addMessageBundlesBuilder(); + messageBundle2.setSourceComputationId(SOURCE_COMPUTATION_ID); + addTimer(workItem2, window(0, 10), new Instant(9), Timer.Type.WATERMARK); + when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(10)); + + runner.processElement(createValue(workItem2, valueCoder)); + runner.finishBundle(); + + runner.startBundle(); + + WorkItem.Builder workItem3 = WorkItem.newBuilder(); + workItem3.setKey(ByteString.copyFromUtf8(KEY)); + workItem3.setWorkToken(WORK_TOKEN); + addTimer(workItem3, window(10, 20), new Instant(19), Timer.Type.WATERMARK); + + when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(20)); + runner.processElement(createDrainingValue(workItem3, valueCoder)); + runner.finishBundle(); + + List>>> result = outputManager.getOutput(outputTag); + + assertThat(result.size(), equalTo(2)); + + assertThat( + result, + containsInAnyOrder( + WindowMatchers.isSingleWindowedValue( + isKv(equalTo(KEY), containsInAnyOrder("v0", "v1", "v2")), + equalTo(window(0, 10).maxTimestamp()), + equalTo(window(0, 10)), + anything(), + is(Boolean.FALSE)), + WindowMatchers.isSingleWindowedValue( + isKv(equalTo(KEY), containsInAnyOrder("v3")), + equalTo(window(10, 20).maxTimestamp()), + equalTo(window(10, 20)), + anything(), + is(Boolean.TRUE)))); + } + @Test public void testSlidingWindows() throws Exception { TupleTag>> outputTag = new TupleTag<>(); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java index c169c9b46a57..52c9844add86 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java @@ -24,6 +24,7 @@ import java.util.Arrays; import java.util.Collection; import java.util.List; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.runners.core.DoFnRunner; import org.apache.beam.runners.core.KeyedWorkItem; import org.apache.beam.runners.core.NullSideInputReader; @@ -114,7 +115,12 @@ private void addElement( valueCoder.encode(value, dataOutput, Context.OUTER); messageBundle .addMessagesBuilder() - .setMetadata(WindmillSink.encodeMetadata(windowsCoder, windows, PaneInfo.NO_FIRING)) + .setMetadata( + WindmillSink.encodeMetadata( + windowsCoder, + windows, + PaneInfo.NO_FIRING, + BeamFnApi.Elements.ElementMetadata.newBuilder().build())) .setData(dataOutput.toByteString()) .setTimestamp(WindmillTimeUtils.harnessToWindmillTimestamp(timestamp)); } @@ -126,7 +132,7 @@ private WindowedValue> createValue( return new ValueInEmptyWindows<>( (KeyedWorkItem) new WindmillKeyedWorkItem<>( - KEY, workItem.build(), windowCoder, wildcardWindowsCoder, valueCoder)); + KEY, workItem.build(), windowCoder, wildcardWindowsCoder, valueCoder, false)); } @Test diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java index e216f912d77f..dd5c36c9cf85 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java @@ -143,6 +143,7 @@ private static Work createMockWork(Windmill.WorkItem workItem, Watermarks waterm watermarks, Work.createProcessingContext( COMPUTATION_ID, new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, Instant::now); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java index ffe71176367a..d66dc24a0c71 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java @@ -22,6 +22,7 @@ import java.io.IOException; import java.util.Collection; import java.util.Collections; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.runners.core.KeyedWorkItem; import org.apache.beam.runners.core.StateNamespace; import org.apache.beam.runners.core.StateNamespaces; @@ -92,7 +93,7 @@ public void testElementIteration() throws Exception { KeyedWorkItem keyedWorkItem = new WindmillKeyedWorkItem<>( - KEY, workItem.build(), WINDOW_CODER, WINDOWS_CODER, VALUE_CODER); + KEY, workItem.build(), WINDOW_CODER, WINDOWS_CODER, VALUE_CODER, null); assertThat( keyedWorkItem.elementsIterable(), @@ -107,10 +108,14 @@ private void addElement( long timestamp, String value, IntervalWindow window, - PaneInfo paneInfo) + PaneInfo pane) throws IOException { ByteString encodedMetadata = - WindmillSink.encodeMetadata(WINDOWS_CODER, Collections.singletonList(window), paneInfo); + WindmillSink.encodeMetadata( + WINDOWS_CODER, + Collections.singletonList(window), + pane, + BeamFnApi.Elements.ElementMetadata.newBuilder().build()); chunk .addMessagesBuilder() .setTimestamp(WindmillTimeUtils.harnessToWindmillTimestamp(new Instant(timestamp))) @@ -143,7 +148,7 @@ public void testTimerOrdering() throws Exception { .build(); KeyedWorkItem keyedWorkItem = - new WindmillKeyedWorkItem<>(KEY, workItem, WINDOW_CODER, WINDOWS_CODER, VALUE_CODER); + new WindmillKeyedWorkItem<>(KEY, workItem, WINDOW_CODER, WINDOWS_CODER, VALUE_CODER, null); assertThat( keyedWorkItem.timersIterable(), diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java index ec8672b6a75f..853fe2ce4c62 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java @@ -96,7 +96,8 @@ public void testTimerDataToFromTimer() { WindmillTimerInternals.windmillTimerToTimerData( prefix, WindmillTimerInternals.timerDataToWindmillTimer(stateFamily, prefix, timer), - coder); + coder, + null); // The function itself bounds output, so we dont expect the original input as the // output, we expect it to be bounded TimerData expected = @@ -145,7 +146,8 @@ public void testTimerDataToFromTimer() { prefix, WindmillTimerInternals.timerDataToWindmillTimer( stateFamily, prefix, timer), - coder), + coder, + null), equalTo(expected)); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java index df3b959c82c5..a3b42cc89c51 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java @@ -204,6 +204,7 @@ private static Work createMockWork(Windmill.WorkItem workItem, Watermarks waterm watermarks, Work.createProcessingContext( COMPUTATION_ID, new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, Instant::now); } @@ -1014,6 +1015,7 @@ public void testFailedWorkItemsAbort() throws Exception { new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, Instant::now); context.start( "key", diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java index c0cb8241d73e..c03aca8bb539 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java @@ -71,6 +71,7 @@ private static ExecutableWork createWork(Windmill.WorkItem workItem) { workItem.getSerializedSize(), Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build(), createWorkProcessingContext(), + null, Instant::now), ignored -> {}); } @@ -82,6 +83,7 @@ private static ExecutableWork expiredWork(Windmill.WorkItem workItem) { workItem.getSerializedSize(), Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build(), createWorkProcessingContext(), + null, () -> Instant.EPOCH), ignored -> {}); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java index 935b25acb6f2..e71a10276051 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java @@ -75,6 +75,7 @@ private static ExecutableWork createWork(ShardedKey shardedKey, long workToken, new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, Instant::now), ignored -> {}); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java index 65e40f171b0c..94c8f4b75957 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java @@ -128,6 +128,7 @@ private static WorkItemScheduler noOpProcessWorkItemFn() { serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies) -> {}; } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSenderTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSenderTest.java index b94270ad7bb7..3217c736adb1 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSenderTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSenderTest.java @@ -68,6 +68,7 @@ public class WindmillStreamSenderTest { serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies) -> {}; @Rule public transient Timeout globalTimeout = Timeout.seconds(600); private ManagedChannel inProcessChannel; diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java index a86e6060955c..a3ea64105e0c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java @@ -83,6 +83,7 @@ private static ExecutableWork createWork(Consumer executeWorkFn) { new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, Instant::now), executeWorkFn); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java index 477c764a70ef..fa1d79d100e2 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java @@ -74,6 +74,7 @@ private static Work createMockWork(long workToken) { throw new UnsupportedOperationException(); }, mock(HeartbeatSender.class)), + null, Instant::now); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java index 5748b128f971..2d319879c79f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java @@ -104,6 +104,7 @@ private static Work createMockWork(long workToken) { throw new UnsupportedOperationException(); }, mock(HeartbeatSender.class)), + null, Instant::now); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStreamTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStreamTest.java index 419000178381..76883bebdac0 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStreamTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStreamTest.java @@ -70,6 +70,7 @@ public class GrpcDirectGetWorkStreamTest { serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies) -> {}; private static final Windmill.JobHeader TEST_JOB_HEADER = Windmill.JobHeader.newBuilder() @@ -283,6 +284,7 @@ public void testConsumedWorkItem_computesAndSendsCorrectExtension() throws Inter serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies) -> { scheduledWorkItems.add(work); }); @@ -327,8 +329,12 @@ public void testConsumedWorkItem_doesNotSendExtensionIfOutstandingBudgetHigh() createGetWorkStream( testStub, initialBudget, - (work, serializedWorkItemSize, watermarks, processingContext, getWorkStreamLatencies) -> - scheduledWorkItems.add(work)); + (work, + serializedWorkItemSize, + watermarks, + processingContext, + drainMode, + getWorkStreamLatencies) -> scheduledWorkItems.add(work)); Windmill.WorkItem workItem = Windmill.WorkItem.newBuilder() .setKey(ByteString.copyFromUtf8("somewhat_long_key")) @@ -365,6 +371,7 @@ public void testConsumedWorkItems() throws InterruptedException { serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies) -> { scheduledWorkItems.add(work); }); @@ -408,6 +415,7 @@ public void testConsumedWorkItems_itemsSplitAcrossResponses() throws Interrupted serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies) -> { scheduledWorkItems.add(work); }); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServerTest.java index e52b6e8de4bf..6caffb485f0b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServerTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServerTest.java @@ -336,6 +336,7 @@ public void onCompleted() { (String computation, @Nullable Instant inputDataWatermark, Instant synchronizedProcessingTime, + Boolean drainMode, WorkItem workItem, long serializedWorkItemSize, ImmutableList getWorkStreamLatencies) -> { @@ -469,6 +470,7 @@ public void onCompleted() { (String computation, @Nullable Instant inputDataWatermark, Instant synchronizedProcessingTime, + Boolean drainMode, WorkItem workItem, long serializedWorkItemSize, ImmutableList getWorkStreamLatencies) -> { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java index f55549f7e2d9..27f20c4cc0f1 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java @@ -95,6 +95,7 @@ private static ExecutableWork createWork(Supplier clock, Consumer new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, clock), processWorkFn); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java index 115deccf6df4..413dd033416f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java @@ -133,6 +133,7 @@ private ExecutableWork createOldWork( Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build(), Work.createProcessingContext( "computationId", new FakeGetDataClient(), ignored -> {}, heartbeatSender), + null, A_LONG_TIME_AGO), processWork); } diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/TimerReceiverFactory.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/TimerReceiverFactory.java index d99a733a4eb4..c2ad60d6280d 100644 --- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/TimerReceiverFactory.java +++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/TimerReceiverFactory.java @@ -89,12 +89,17 @@ public FnDataReceiver> create(String transformId, String timerFamil StateNamespace namespace = StateNamespaces.window(windowCoder, (BoundedWindow) window); TimerInternals.TimerData timerData = TimerInternals.TimerData.of( - timer.getDynamicTimerTag(), - encodeToTimerDataTimerId(timerSpec.transformId(), timerSpec.timerId()), - namespace, - timer.getClearBit() ? BoundedWindow.TIMESTAMP_MAX_VALUE : timer.getFireTimestamp(), - timer.getClearBit() ? BoundedWindow.TIMESTAMP_MAX_VALUE : timer.getHoldTimestamp(), - timerSpec.getTimerSpec().getTimeDomain()); + timer.getDynamicTimerTag(), + encodeToTimerDataTimerId(timerSpec.transformId(), timerSpec.timerId()), + namespace, + timer.getClearBit() + ? BoundedWindow.TIMESTAMP_MAX_VALUE + : timer.getFireTimestamp(), + timer.getClearBit() + ? BoundedWindow.TIMESTAMP_MAX_VALUE + : timer.getHoldTimestamp(), + timerSpec.getTimerSpec().getTimeDomain()) + .draining(timer.getDraining()); timerDataConsumer.accept(timer, timerData); } }; diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/translation/PipelineTranslatorUtils.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/translation/PipelineTranslatorUtils.java index 5b7880ad56ee..3848bd99c127 100644 --- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/translation/PipelineTranslatorUtils.java +++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/translation/PipelineTranslatorUtils.java @@ -156,7 +156,8 @@ private static void fireTimer( Collections.singletonList(window), timestamp, outputTimestamp, - PaneInfo.NO_FIRING); + PaneInfo.NO_FIRING, + null); KV transformAndTimerId = TimerReceiverFactory.decodeTimerDataTimerId(timer.getTimerFamilyId()); FnDataReceiver fnTimerReceiver = timerReceivers.get(transformAndTimerId); diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java index 556cc7993ae1..d6477095bf96 100644 --- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java +++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java @@ -2254,6 +2254,7 @@ private org.apache.beam.sdk.util.construction.Timer timerForTest( Collections.singletonList(GlobalWindow.INSTANCE), BoundedWindow.TIMESTAMP_MIN_VALUE.plus(Duration.millis(fireTimestamp)), BoundedWindow.TIMESTAMP_MIN_VALUE.plus(Duration.millis(holdTimestamp)), - PaneInfo.NO_FIRING); + PaneInfo.NO_FIRING, + null); } } diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/wire/CommonCoderTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/wire/CommonCoderTest.java index eccf1e66434e..839dc898d3c1 100644 --- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/wire/CommonCoderTest.java +++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/wire/CommonCoderTest.java @@ -326,7 +326,8 @@ private static Object convertValue(Object value, CommonCoder coderSpec, Coder co windows, new Instant(((Number) kvMap.get("fireTimestamp")).longValue()), new Instant(((Number) kvMap.get("holdTimestamp")).longValue()), - paneInfo); + paneInfo, + null); } else if (s.equals(getUrn(StandardCoders.Enum.INTERVAL_WINDOW))) { Map kvMap = (Map) value; Instant end = new Instant(((Number) kvMap.get("end")).longValue()); diff --git a/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/StatefulParDoP.java b/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/StatefulParDoP.java index f696873ec5c0..24f298c7e783 100644 --- a/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/StatefulParDoP.java +++ b/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/StatefulParDoP.java @@ -105,7 +105,8 @@ private static void fireTimer( window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain()); + timer.getDomain(), + null); } @Override 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 index b815649a7659..237b9ba19a42 100644 --- 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 @@ -22,6 +22,7 @@ import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -61,11 +62,19 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { withMetrics( () -> underlying.onTimer( - timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain), + timerId, + timerFamilyId, + key, + window, + timestamp, + outputTimestamp, + timeDomain, + draining), false); } 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 index d07a9bda78c6..08f4b367e9c2 100644 --- 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 @@ -154,8 +154,10 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { - underlying.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + TimeDomain timeDomain, + @Nullable Boolean draining) { + underlying.onTimer( + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } @Override 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 index 2f27e31d05dd..45c2ec4bde89 100644 --- 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 @@ -437,7 +437,8 @@ private void fireTimer(KeyedTimerData keyedTimerData) { window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain()); + timer.getDomain(), + null); } // todo: should this go through bundle manager to start and finish the bundle? diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnRunnerWithKeyedInternals.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnRunnerWithKeyedInternals.java index 34e3405660c2..8daefb3e3da2 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnRunnerWithKeyedInternals.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnRunnerWithKeyedInternals.java @@ -24,6 +24,7 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow; 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; /** This class wraps a DoFnRunner with keyed StateInternals and TimerInternals access. */ @@ -66,14 +67,15 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { // Note: wrap with KV.of(key, null) as a special use case of setKeyedInternals() to set key // directly. setKeyedInternals(KV.of(key, null)); try { underlying.onTimer( - timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } finally { clearKeyedInternals(); } diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java index 743a42d14791..b7f6726f81e8 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java @@ -425,7 +425,8 @@ private void fireTimer(KeyedTimerData keyedTimerData) { window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain()); + timer.getDomain(), + null); } // todo: should this go through bundle manager to start and finish the bundle? diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java index 7129ba9145eb..d012b33b428f 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java @@ -69,6 +69,7 @@ import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.samza.context.Context; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** A factory for Samza runner translator to create underlying DoFnRunner used in {@link DoFnOp}. */ @@ -430,7 +431,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { final KV timerReceiverKey = TimerReceiverFactory.decodeTimerDataTimerId(timerFamilyId); final FnDataReceiver timerReceiver = @@ -443,7 +445,8 @@ public void onTimer( timestamp, outputTimestamp, // TODO: Support propagating the PaneInfo through. - PaneInfo.NO_FIRING); + PaneInfo.NO_FIRING, + draining); try { timerReceiver.accept(timerValue); } catch (Exception e) { diff --git a/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerFactory.java b/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerFactory.java index 350f7daa56cd..3bbc11f84f3f 100644 --- a/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerFactory.java +++ b/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerFactory.java @@ -47,6 +47,7 @@ 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.collect.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; import scala.Serializable; @@ -266,7 +267,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { throw new UnsupportedOperationException(); } diff --git a/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java b/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java index db4ee5be5786..e9b42b169009 100644 --- a/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java +++ b/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java @@ -29,6 +29,7 @@ import org.apache.beam.sdk.transforms.reflect.DoFnInvokers; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** DoFnRunner decorator which registers {@link MetricsContainer}. */ @@ -77,9 +78,11 @@ public void onTimer( final BoundedWindow window, final Instant timestamp, final Instant outputTimestamp, - final TimeDomain timeDomain) { + final TimeDomain timeDomain, + @Nullable Boolean draining) { try (Closeable ignored = MetricsEnvironment.scopedMetricsContainer(metrics)) { - delegate.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + delegate.onTimer( + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } catch (IOException e) { throw new RuntimeException(e); } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/AbstractInOutIterator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/AbstractInOutIterator.java index 9d73a605b3b5..6008ceb47bcc 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/AbstractInOutIterator.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/AbstractInOutIterator.java @@ -75,7 +75,8 @@ public void fireTimer(TimerInternals.TimerData timer) { window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain()); + timer.getDomain(), + null); } finally { if (this.ctx.getTimerDataIterator() instanceof ParDoStateUpdateFn.SparkTimerInternalsIterator) { diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java index a6b1f65571dc..c26d9f75e98b 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java @@ -28,6 +28,7 @@ import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** DoFnRunner decorator which registers {@link MetricsContainerImpl}. */ @@ -76,9 +77,11 @@ public void onTimer( final BoundedWindow window, final Instant timestamp, final Instant outputTimestamp, - final TimeDomain timeDomain) { + final TimeDomain timeDomain, + @Nullable Boolean draining) { try (Closeable ignored = MetricsEnvironment.scopedMetricsContainer(metricsContainer())) { - delegate.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + delegate.onTimer( + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } catch (IOException e) { throw new RuntimeException(e); } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java index 03735355de51..0760c1aeb649 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java @@ -115,6 +115,11 @@ public PaneInfo getPaneInfo() { return null; } + @Override + public @Nullable Boolean isDraining() { + return null; + } + @Override public @Nullable Long getRecordOffset() { return null; diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/AbstractInOutIteratorTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/AbstractInOutIteratorTest.java index b1e0e3a39808..e6da459e7c71 100644 --- a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/AbstractInOutIteratorTest.java +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/AbstractInOutIteratorTest.java @@ -107,7 +107,8 @@ public void testFireTimer() { mockWindow, TEST_TIMESTAMP, TEST_OUTPUT_TIMESTAMP, - TEST_TIME_DOMAIN); + TEST_TIME_DOMAIN, + null); // Verify that timer data iterator deletion was not called (no timer iterator was set in this // test) @@ -133,7 +134,8 @@ public void testFireTimerWithTimerDataIterator() { mockWindow, TEST_TIMESTAMP, TEST_OUTPUT_TIMESTAMP, - TEST_TIME_DOMAIN); + TEST_TIME_DOMAIN, + null); // Verify that the timer data iterator's deleteTimer method was called verify(mockTimerDataIterator).deleteTimer(mockTimer); diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkInputDataProcessorTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkInputDataProcessorTest.java index 2ff06b59f8eb..346587993472 100644 --- a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkInputDataProcessorTest.java +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkInputDataProcessorTest.java @@ -40,6 +40,7 @@ import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterators; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; import org.junit.Rule; import org.junit.Test; @@ -252,7 +253,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) {} + TimeDomain timeDomain, + @Nullable Boolean draining) {} @Override public void finishBundle() {} diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java index d0714de60328..59bcde70edf1 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java @@ -361,6 +361,9 @@ public abstract class ProcessContext extends WindowedContext { @Pure public abstract Long currentRecordOffset(); + + @Pure + public abstract Boolean draining(); } /** Information accessible when running a {@link DoFn.OnTimer} method. */ @@ -377,6 +380,9 @@ public abstract class OnTimerContext extends WindowedContext { /** Returns the time domain of the current timer. */ public abstract TimeDomain timeDomain(); + + @Pure + public abstract Boolean draining(); } public abstract class OnWindowExpirationContext extends WindowedContext { @@ -858,6 +864,15 @@ public interface MultiOutputReceiver { @Target(ElementType.PARAMETER) public @interface Element {} + /** + * Parameter annotation for the drain information for {@link ProcessElement} and {@link OnTimer} + * methods. + */ + @Documented + @Retention(RetentionPolicy.RUNTIME) + @Target(ElementType.PARAMETER) + public @interface Draining {} + /** * Parameter annotation for the restriction for {@link GetSize}, {@link SplitRestriction}, {@link * GetInitialWatermarkEstimatorState}, {@link NewWatermarkEstimator}, and {@link NewTracker} diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java index c59d6b528c3f..2269f1fbd799 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java @@ -490,7 +490,7 @@ public void output(TupleTag tag, T output, Instant timestamp, BoundedWind getMutableOutput(tag) .add( ValueInSingleWindow.of( - output, timestamp, window, PaneInfo.NO_FIRING, null, null)); + output, timestamp, window, PaneInfo.NO_FIRING, null, null, null)); } @Override @@ -519,7 +519,8 @@ public void output( window, PaneInfo.NO_FIRING, currentRecordId, - currentRecordOffset)); + currentRecordOffset, + null)); } }; } @@ -608,6 +609,11 @@ public PaneInfo pane() { return element.getPaneInfo(); } + @Override + public Boolean draining() { + return element.isDraining(); + } + @Override public String currentRecordId() { return element.getCurrentRecordId(); @@ -670,7 +676,7 @@ public void outputWithTimestamp(TupleTag tag, T output, Instant timestamp getMutableOutput(tag) .add( ValueInSingleWindow.of( - output, timestamp, element.getWindow(), element.getPaneInfo(), null, null)); + output, timestamp, element.getWindow(), element.getPaneInfo(), null, null, null)); } @Override @@ -682,7 +688,7 @@ public void outputWindowedValue( PaneInfo paneInfo) { for (BoundedWindow w : windows) { getMutableOutput(tag) - .add(ValueInSingleWindow.of(output, timestamp, w, paneInfo, null, null)); + .add(ValueInSingleWindow.of(output, timestamp, w, paneInfo, null, null, null)); } } @@ -699,7 +705,7 @@ public void outputWindowedValue( getMutableOutput(tag) .add( ValueInSingleWindow.of( - output, timestamp, w, paneInfo, currentRecordId, currentRecordOffset)); + output, timestamp, w, paneInfo, currentRecordId, currentRecordOffset, null)); } } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java index a01b5f570a57..7382f29310c5 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java @@ -179,6 +179,7 @@ public Duration getAllowedTimestampSkew() { public void processElement( @Element KV> kv, OutputReceiver> outputReceiver) { + // todo #33176 specify additional metadata in the future outputReceiver .builder(KV.of(kv.getKey(), kv.getValue().getValue())) .setTimestamp(kv.getValue().getTimestamp()) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reify.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reify.java index 797af9538c53..af125d9e63e8 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reify.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reify.java @@ -136,6 +136,7 @@ public PCollection>> expand(PCollection> i KvCoder coder = (KvCoder) input.getCoder(); return input .apply( + // todo #33176 specify additional metadata in the future ParDo.of( new DoFn, KV>>() { @ProcessElement diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java index b2de48342d7c..0a8d058107b8 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java @@ -184,6 +184,7 @@ public Duration getAllowedTimestampSkew() { public void processElement( @Element KV> kv, OutputReceiver> outputReceiver) { + // todo #33176 specify additional metadata in the future outputReceiver .builder(KV.of(kv.getKey(), kv.getValue().getValue())) .setTimestamp(kv.getValue().getTimestamp()) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java index 6e4c694d48e3..bc83687bae4e 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java @@ -146,10 +146,10 @@ private static byte encodedByte(boolean isFirst, boolean isLast, Timing timing) ImmutableMap.Builder decodingBuilder = ImmutableMap.builder(); for (Timing timing : Timing.values()) { long onTimeIndex = timing == Timing.EARLY ? -1 : 0; - register(decodingBuilder, new PaneInfo(true, true, timing, 0, onTimeIndex)); - register(decodingBuilder, new PaneInfo(true, false, timing, 0, onTimeIndex)); - register(decodingBuilder, new PaneInfo(false, true, timing, -1, onTimeIndex)); - register(decodingBuilder, new PaneInfo(false, false, timing, -1, onTimeIndex)); + register(decodingBuilder, new PaneInfo(true, true, timing, 0, onTimeIndex, false)); + register(decodingBuilder, new PaneInfo(true, false, timing, 0, onTimeIndex, false)); + register(decodingBuilder, new PaneInfo(false, true, timing, -1, onTimeIndex, false)); + register(decodingBuilder, new PaneInfo(false, false, timing, -1, onTimeIndex, false)); } BYTE_TO_PANE_INFO = decodingBuilder.build(); } @@ -159,7 +159,7 @@ private static void register(ImmutableMap.Builder builder, PaneI } private final byte encodedByte; - + private final boolean containsElementMetadata; private final boolean isFirst; private final boolean isLast; private final Timing timing; @@ -177,13 +177,20 @@ private static void register(ImmutableMap.Builder builder, PaneI public static final PaneInfo ON_TIME_AND_ONLY_FIRING = PaneInfo.createPane(true, true, Timing.ON_TIME, 0, 0); - private PaneInfo(boolean isFirst, boolean isLast, Timing timing, long index, long onTimeIndex) { + private PaneInfo( + boolean isFirst, + boolean isLast, + Timing timing, + long index, + long onTimeIndex, + boolean containsElementMetadata) { this.encodedByte = encodedByte(isFirst, isLast, timing); this.isFirst = isFirst; this.isLast = isLast; this.timing = timing; this.index = index; this.nonSpeculativeIndex = onTimeIndex; + this.containsElementMetadata = containsElementMetadata; } public static PaneInfo createPane(boolean isFirst, boolean isLast, Timing timing) { @@ -194,10 +201,21 @@ public static PaneInfo createPane(boolean isFirst, boolean isLast, Timing timing /** Factory method to create a {@link PaneInfo} with the specified parameters. */ public static PaneInfo createPane( boolean isFirst, boolean isLast, Timing timing, long index, long onTimeIndex) { + return createPane(isFirst, isLast, timing, index, onTimeIndex, false); + } + + /** Factory method to create a {@link PaneInfo} with the specified parameters. */ + public static PaneInfo createPane( + boolean isFirst, + boolean isLast, + Timing timing, + long index, + long onTimeIndex, + boolean containsElementMetadata) { if (isFirst || timing == Timing.UNKNOWN) { return checkNotNull(BYTE_TO_PANE_INFO.get(encodedByte(isFirst, isLast, timing))); } else { - return new PaneInfo(isFirst, isLast, timing, index, onTimeIndex); + return new PaneInfo(isFirst, isLast, timing, index, onTimeIndex, containsElementMetadata); } } @@ -219,6 +237,15 @@ public boolean isFirst() { return isFirst; } + public boolean isElementMetadata() { + return containsElementMetadata; + } + + public PaneInfo withElementMetadata(boolean elementMetadata) { + return new PaneInfo( + this.isFirst, this.isLast, this.timing, index, nonSpeculativeIndex, elementMetadata); + } + /** Return true if this is the last pane that will be produced in the associated window. */ public boolean isLast() { return isLast; @@ -295,6 +322,8 @@ public String toString() { /** A Coder for encoding PaneInfo instances. */ public static class PaneInfoCoder extends AtomicCoder { + private static final byte ELEMENT_METADATA_MASK = (byte) 0x80; + private enum Encoding { FIRST, ONE_INDEX, @@ -337,16 +366,17 @@ private PaneInfoCoder() {} public void encode(PaneInfo value, final OutputStream outStream) throws CoderException, IOException { Encoding encoding = chooseEncoding(value); + byte elementMetadata = value.containsElementMetadata ? ELEMENT_METADATA_MASK : 0x00; switch (chooseEncoding(value)) { case FIRST: - outStream.write(value.encodedByte); + outStream.write(value.encodedByte | elementMetadata); break; case ONE_INDEX: - outStream.write(value.encodedByte | encoding.tag); + outStream.write(value.encodedByte | encoding.tag | elementMetadata); VarInt.encode(value.index, outStream); break; case TWO_INDICES: - outStream.write(value.encodedByte | encoding.tag); + outStream.write(value.encodedByte | encoding.tag | elementMetadata); VarInt.encode(value.index, outStream); VarInt.encode(value.nonSpeculativeIndex, outStream); break; @@ -360,9 +390,10 @@ public PaneInfo decode(final InputStream inStream) throws CoderException, IOExce byte keyAndTag = (byte) inStream.read(); PaneInfo base = Preconditions.checkNotNull(BYTE_TO_PANE_INFO.get((byte) (keyAndTag & 0x0F))); long index, onTimeIndex; - switch (Encoding.fromTag(keyAndTag)) { + boolean elementMetadata = (keyAndTag & ELEMENT_METADATA_MASK) != 0; + switch (Encoding.fromTag((byte) (keyAndTag & ~ELEMENT_METADATA_MASK))) { case FIRST: - return base; + return base.withElementMetadata(elementMetadata); case ONE_INDEX: index = VarInt.decodeLong(inStream); onTimeIndex = base.timing == Timing.EARLY ? -1 : index; @@ -374,7 +405,8 @@ public PaneInfo decode(final InputStream inStream) throws CoderException, IOExce default: throw new CoderException("Unknown encoding " + (keyAndTag & 0xF0)); } - return new PaneInfo(base.isFirst, base.isLast, base.timing, index, onTimeIndex); + return new PaneInfo( + base.isFirst, base.isLast, base.timing, index, onTimeIndex, elementMetadata); } @Override diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/SplittableParDoNaiveBounded.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/SplittableParDoNaiveBounded.java index e6394b8810a4..b4ce5594f382 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/SplittableParDoNaiveBounded.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/SplittableParDoNaiveBounded.java @@ -536,6 +536,11 @@ public Object key() { throw new UnsupportedOperationException(); } + @Override + public Boolean draining() { + return outerContext.draining(); + } + @Override public Object sideInput(String tagId) { PCollectionView view = sideInputMapping.get(tagId); diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Timer.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Timer.java index d443f008b7d2..52cb87f96fa7 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Timer.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Timer.java @@ -65,9 +65,10 @@ public static Timer of( Collection windows, Instant fireTimestamp, Instant holdTimestamp, - PaneInfo paneInfo) { + PaneInfo paneInfo, + Boolean draining) { return new AutoValue_Timer( - userKey, dynamicTimerTag, windows, false, fireTimestamp, holdTimestamp, paneInfo); + userKey, dynamicTimerTag, windows, false, fireTimestamp, holdTimestamp, paneInfo, draining); } /** @@ -76,7 +77,7 @@ public static Timer of( */ public static Timer cleared( K userKey, String dynamicTimerTag, Collection windows) { - return new AutoValue_Timer(userKey, dynamicTimerTag, windows, true, null, null, null); + return new AutoValue_Timer(userKey, dynamicTimerTag, windows, true, null, null, null, false); } /** Returns the key that the timer is set on. */ @@ -116,6 +117,8 @@ public static Timer cleared( */ public abstract @Nullable PaneInfo getPaneInfo(); + public abstract @Nullable Boolean getDraining(); + @Override public final boolean equals(@Nullable Object other) { if (!(other instanceof Timer)) { @@ -201,7 +204,8 @@ public Timer decode(InputStream inStream) throws IOException { Instant fireTimestamp = InstantCoder.of().decode(inStream); Instant holdTimestamp = InstantCoder.of().decode(inStream); PaneInfo paneInfo = PaneInfoCoder.INSTANCE.decode(inStream); - return Timer.of(userKey, dynamicTimerTag, windows, fireTimestamp, holdTimestamp, paneInfo); + return Timer.of( + userKey, dynamicTimerTag, windows, fireTimestamp, holdTimestamp, paneInfo, false); } @Override diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/OutputBuilder.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/OutputBuilder.java index a7f8bc8e03b1..5762d32ae832 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/OutputBuilder.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/OutputBuilder.java @@ -48,5 +48,7 @@ public interface OutputBuilder extends WindowedValue { OutputBuilder setRecordOffset(@Nullable Long recordOffset); + OutputBuilder setDraining(@Nullable Boolean drain); + void output(); } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java index 7dc5fef52ecb..74108328baf0 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java @@ -22,7 +22,9 @@ import java.io.InputStream; import java.io.OutputStream; import java.util.List; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.coders.ByteArrayCoder; import org.apache.beam.sdk.coders.InstantCoder; import org.apache.beam.sdk.coders.StructuredCoder; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; @@ -64,20 +66,24 @@ public T getValue() { public abstract @Nullable Long getCurrentRecordOffset(); + public abstract @Nullable Boolean isDraining(); + + // todo #33176 specify additional metadata in the future public static ValueInSingleWindow of( T value, Instant timestamp, BoundedWindow window, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { return new AutoValue_ValueInSingleWindow<>( - value, timestamp, window, paneInfo, currentRecordId, currentRecordOffset); + value, timestamp, window, paneInfo, currentRecordId, currentRecordOffset, draining); } public static ValueInSingleWindow of( T value, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) { - return of(value, timestamp, window, paneInfo, null, null); + return of(value, timestamp, window, paneInfo, null, null, null); } /** A coder for {@link ValueInSingleWindow}. */ @@ -110,7 +116,26 @@ public void encode(ValueInSingleWindow windowedElem, OutputStream outStream, throws IOException { InstantCoder.of().encode(windowedElem.getTimestamp(), outStream); windowCoder.encode(windowedElem.getWindow(), outStream); - PaneInfo.PaneInfoCoder.INSTANCE.encode(windowedElem.getPaneInfo(), outStream); + boolean metadataSupported = WindowedValues.WindowedValueCoder.isMetadataSupported(); + PaneInfo.PaneInfoCoder.INSTANCE.encode( + windowedElem.getPaneInfo().withElementMetadata(metadataSupported), outStream); + if (metadataSupported) { + BeamFnApi.Elements.ElementMetadata.Builder builder = + BeamFnApi.Elements.ElementMetadata.newBuilder(); + // todo #33176 specify additional metadata in the future + Boolean draining = windowedElem.isDraining(); + if (draining != null) { + builder.setDrain( + draining + ? BeamFnApi.Elements.DrainMode.Enum.DRAINING + : BeamFnApi.Elements.DrainMode.Enum.NOT_DRAINING); + } else { + builder.setDrain(BeamFnApi.Elements.DrainMode.Enum.UNSPECIFIED); + } + BeamFnApi.Elements.ElementMetadata metadata = builder.build(); + ByteArrayCoder.of().encode(metadata.toByteArray(), outStream); + } + valueCoder.encode(windowedElem.getValue(), outStream, context); } @@ -124,8 +149,20 @@ public ValueInSingleWindow decode(InputStream inStream, Context context) thro Instant timestamp = InstantCoder.of().decode(inStream); BoundedWindow window = windowCoder.decode(inStream); PaneInfo paneInfo = PaneInfo.PaneInfoCoder.INSTANCE.decode(inStream); + Boolean draining = null; + if (WindowedValues.WindowedValueCoder.isMetadataSupported() && paneInfo.isElementMetadata()) { + BeamFnApi.Elements.ElementMetadata elementMetadata = + BeamFnApi.Elements.ElementMetadata.parseFrom(ByteArrayCoder.of().decode(inStream)); + if (elementMetadata.hasDrain() + && elementMetadata.getDrain() != BeamFnApi.Elements.DrainMode.Enum.UNSPECIFIED) { + draining = elementMetadata.getDrain() == BeamFnApi.Elements.DrainMode.Enum.DRAINING; + } + } + T value = valueCoder.decode(inStream, context); - return new AutoValue_ValueInSingleWindow<>(value, timestamp, window, paneInfo, null, null); + // todo #33176 specify additional metadata in the future + return new AutoValue_ValueInSingleWindow<>( + value, timestamp, window, paneInfo, null, null, draining); } @Override diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java index ea6be129ecb4..762a602bc3f9 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java @@ -52,6 +52,9 @@ public interface WindowedValue { @Nullable Long getRecordOffset(); + @Nullable + Boolean isDraining(); + /** * A representation of each of the actual values represented by this compressed {@link * WindowedValue}, one per window. diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java index 9b079b8699b9..1f458a4a7729 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java @@ -36,6 +36,7 @@ import java.util.List; import java.util.Objects; import java.util.Set; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.sdk.coders.ByteArrayCoder; import org.apache.beam.sdk.coders.Coder; @@ -98,6 +99,7 @@ public static class Builder implements OutputBuilder { private @MonotonicNonNull Collection windows; private @Nullable String recordId; private @Nullable Long recordOffset; + private @Nullable Boolean draining; @Override public Builder setValue(T value) { @@ -141,6 +143,12 @@ public Builder setRecordOffset(@Nullable Long recordOffset) { return this; } + @Override + public Builder setDraining(@Nullable Boolean draining) { + this.draining = draining; + return this; + } + public Builder setReceiver(WindowedValueReceiver receiver) { this.receiver = receiver; return this; @@ -189,6 +197,11 @@ public PaneInfo getPaneInfo() { return recordOffset; } + @Override + public @Nullable Boolean isDraining() { + return draining; + } + @Override public Collection> explodeWindows() { throw new UnsupportedOperationException( @@ -217,7 +230,8 @@ public void output() { } public WindowedValue build() { - return WindowedValues.of(getValue(), getTimestamp(), getWindows(), getPaneInfo()); + return WindowedValues.of( + getValue(), getTimestamp(), getWindows(), getPaneInfo(), null, null, isDraining()); } @Override @@ -227,6 +241,7 @@ public String toString() { .add("timestamp", getTimestamp()) .add("windows", getWindows()) .add("paneInfo", getPaneInfo()) + .add("draining", isDraining()) .add("receiver", receiver) .toString(); } @@ -234,7 +249,7 @@ public String toString() { public static WindowedValue of( T value, Instant timestamp, Collection windows, PaneInfo paneInfo) { - return of(value, timestamp, windows, paneInfo, null, null); + return of(value, timestamp, windows, paneInfo, null, null, null); } /** Returns a {@code WindowedValue} with the given value, timestamp, and windows. */ @@ -244,27 +259,32 @@ public static WindowedValue of( Collection windows, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { checkArgument(paneInfo != null, "WindowedValue requires PaneInfo, but it was null"); checkArgument(windows.size() > 0, "WindowedValue requires windows, but there were none"); if (windows.size() == 1) { - return of(value, timestamp, windows.iterator().next(), paneInfo); + return of(value, timestamp, windows.iterator().next(), paneInfo, draining); } else { return new TimestampedValueInMultipleWindows<>( - value, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset); + value, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset, draining); } } /** @deprecated for use only in compatibility with old broken code */ @Deprecated static WindowedValue createWithoutValidation( - T value, Instant timestamp, Collection windows, PaneInfo paneInfo) { + T value, + Instant timestamp, + Collection windows, + PaneInfo paneInfo, + @Nullable Boolean draining) { if (windows.size() == 1) { - return of(value, timestamp, windows.iterator().next(), paneInfo); + return of(value, timestamp, windows.iterator().next(), paneInfo, draining); } else { return new TimestampedValueInMultipleWindows<>( - value, timestamp, windows, paneInfo, null, null); + value, timestamp, windows, paneInfo, null, null, draining); } } @@ -273,13 +293,26 @@ public static WindowedValue of( T value, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) { checkArgument(paneInfo != null, "WindowedValue requires PaneInfo, but it was null"); + return of(value, timestamp, window, paneInfo, null); + } + + /** Returns a {@code WindowedValue} with the given value, timestamp, and window. */ + public static WindowedValue of( + T value, + Instant timestamp, + BoundedWindow window, + PaneInfo paneInfo, + @Nullable Boolean draining) { + checkArgument(paneInfo != null, "WindowedValue requires PaneInfo, but it was null"); + boolean isGlobal = GlobalWindow.INSTANCE.equals(window); if (isGlobal && BoundedWindow.TIMESTAMP_MIN_VALUE.equals(timestamp)) { return valueInGlobalWindow(value, paneInfo); } else if (isGlobal) { - return new TimestampedValueInGlobalWindow<>(value, timestamp, paneInfo, null, null); + return new TimestampedValueInGlobalWindow<>(value, timestamp, paneInfo, null, null, draining); } else { - return new TimestampedValueInSingleWindow<>(value, timestamp, window, paneInfo, null, null); + return new TimestampedValueInSingleWindow<>( + value, timestamp, window, paneInfo, null, null, draining); } } @@ -288,7 +321,7 @@ public static WindowedValue of( * default timestamp and pane. */ public static WindowedValue valueInGlobalWindow(T value) { - return new ValueInGlobalWindow<>(value, PaneInfo.NO_FIRING, null, null); + return new ValueInGlobalWindow<>(value, PaneInfo.NO_FIRING, null, null, null); } /** @@ -296,7 +329,7 @@ public static WindowedValue valueInGlobalWindow(T value) { * default timestamp and the specified pane. */ public static WindowedValue valueInGlobalWindow(T value, PaneInfo paneInfo) { - return new ValueInGlobalWindow<>(value, paneInfo, null, null); + return new ValueInGlobalWindow<>(value, paneInfo, null, null, null); } /** @@ -307,7 +340,8 @@ public static WindowedValue timestampedValueInGlobalWindow(T value, Insta if (BoundedWindow.TIMESTAMP_MIN_VALUE.equals(timestamp)) { return valueInGlobalWindow(value); } else { - return new TimestampedValueInGlobalWindow<>(value, timestamp, PaneInfo.NO_FIRING, null, null); + return new TimestampedValueInGlobalWindow<>( + value, timestamp, PaneInfo.NO_FIRING, null, null, null); } } @@ -320,7 +354,7 @@ public static WindowedValue timestampedValueInGlobalWindow( if (paneInfo.equals(PaneInfo.NO_FIRING)) { return timestampedValueInGlobalWindow(value, timestamp); } else { - return new TimestampedValueInGlobalWindow<>(value, timestamp, paneInfo, null, null); + return new TimestampedValueInGlobalWindow<>(value, timestamp, paneInfo, null, null, null); } } @@ -336,7 +370,8 @@ public static WindowedValue withValue( windowedValue.getWindows(), windowedValue.getPaneInfo(), windowedValue.getRecordId(), - windowedValue.getRecordOffset()); + windowedValue.getRecordOffset(), + windowedValue.isDraining()); } public static boolean equals( @@ -387,6 +422,7 @@ private abstract static class SimpleWindowedValue implements WindowedValue private final PaneInfo paneInfo; private final @Nullable String currentRecordId; private final @Nullable Long currentRecordOffset; + private final @Nullable Boolean draining; @Override public @Nullable String getRecordId() { @@ -398,15 +434,22 @@ private abstract static class SimpleWindowedValue implements WindowedValue return currentRecordOffset; } + @Override + public @Nullable Boolean isDraining() { + return draining; + } + protected SimpleWindowedValue( T value, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { this.value = value; this.paneInfo = checkNotNull(paneInfo); this.currentRecordId = currentRecordId; this.currentRecordOffset = currentRecordOffset; + this.draining = draining; } @Override @@ -454,8 +497,9 @@ public MinTimestampWindowedValue( T value, PaneInfo pane, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { - super(value, pane, currentRecordId, currentRecordOffset); + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { + super(value, pane, currentRecordId, currentRecordOffset, draining); } @Override @@ -472,8 +516,9 @@ public ValueInGlobalWindow( T value, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { - super(value, paneInfo, currentRecordId, currentRecordOffset); + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { + super(value, paneInfo, currentRecordId, currentRecordOffset, draining); } @Override @@ -488,7 +533,8 @@ public BoundedWindow getWindow() { @Override public WindowedValue withValue(NewT newValue) { - return new ValueInGlobalWindow<>(newValue, getPaneInfo(), getRecordId(), getRecordOffset()); + return new ValueInGlobalWindow<>( + newValue, getPaneInfo(), getRecordId(), getRecordOffset(), isDraining()); } @Override @@ -512,6 +558,7 @@ public String toString() { return MoreObjects.toStringHelper(getClass()) .add("value", getValue()) .add("paneInfo", getPaneInfo()) + .add("draining", isDraining()) .toString(); } } @@ -525,8 +572,9 @@ public TimestampedWindowedValue( Instant timestamp, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { - super(value, paneInfo, currentRecordId, currentRecordOffset); + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { + super(value, paneInfo, currentRecordId, currentRecordOffset, draining); this.timestamp = checkNotNull(timestamp); } @@ -548,8 +596,9 @@ public TimestampedValueInGlobalWindow( Instant timestamp, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { - super(value, timestamp, paneInfo, currentRecordId, currentRecordOffset); + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { + super(value, timestamp, paneInfo, currentRecordId, currentRecordOffset, draining); } @Override @@ -565,7 +614,7 @@ public BoundedWindow getWindow() { @Override public WindowedValue withValue(NewT newValue) { return new TimestampedValueInGlobalWindow<>( - newValue, getTimestamp(), getPaneInfo(), getRecordId(), getRecordOffset()); + newValue, getTimestamp(), getPaneInfo(), getRecordId(), getRecordOffset(), isDraining()); } @Override @@ -595,6 +644,7 @@ public String toString() { .add("value", getValue()) .add("timestamp", getTimestamp()) .add("paneInfo", getPaneInfo()) + .add("draining", isDraining()) .toString(); } } @@ -614,15 +664,22 @@ public TimestampedValueInSingleWindow( BoundedWindow window, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { - super(value, timestamp, paneInfo, currentRecordId, currentRecordOffset); + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { + super(value, timestamp, paneInfo, currentRecordId, currentRecordOffset, draining); this.window = checkNotNull(window); } @Override public WindowedValue withValue(NewT newValue) { return new TimestampedValueInSingleWindow<>( - newValue, getTimestamp(), window, getPaneInfo(), getRecordId(), getRecordOffset()); + newValue, + getTimestamp(), + window, + getPaneInfo(), + getRecordId(), + getRecordOffset(), + isDraining()); } @Override @@ -664,6 +721,7 @@ public String toString() { .add("timestamp", getTimestamp()) .add("window", window) .add("paneInfo", getPaneInfo()) + .add("draining", isDraining()) .toString(); } } @@ -678,8 +736,9 @@ public TimestampedValueInMultipleWindows( Collection windows, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { - super(value, timestamp, paneInfo, currentRecordId, currentRecordOffset); + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { + super(value, timestamp, paneInfo, currentRecordId, currentRecordOffset, draining); this.windows = checkNotNull(windows); } @@ -691,7 +750,13 @@ public Collection getWindows() { @Override public WindowedValue withValue(NewT newValue) { return new TimestampedValueInMultipleWindows<>( - newValue, getTimestamp(), getWindows(), getPaneInfo(), getRecordId(), getRecordOffset()); + newValue, + getTimestamp(), + getWindows(), + getPaneInfo(), + getRecordId(), + getRecordOffset(), + isDraining()); } @Override @@ -729,6 +794,7 @@ public String toString() { .add("timestamp", getTimestamp()) .add("windows", windows) .add("paneInfo", getPaneInfo()) + .add("draining", isDraining()) .toString(); } @@ -763,6 +829,15 @@ public static ParamWindowedValueCoder getParamWindowedValueCoder(Coder /** Abstract class for {@code WindowedValue} coder. */ public abstract static class WindowedValueCoder extends StructuredCoder> { final Coder valueCoder; + private static boolean metadataSupported = false; + + public static void setMetadataSupported() { + metadataSupported = true; + } + + public static boolean isMetadataSupported() { + return metadataSupported; + } WindowedValueCoder(Coder valueCoder) { this.valueCoder = checkNotNull(valueCoder); @@ -829,7 +904,24 @@ public void encode(WindowedValue windowedElem, OutputStream outStream, Contex throws CoderException, IOException { InstantCoder.of().encode(windowedElem.getTimestamp(), outStream); windowsCoder.encode(windowedElem.getWindows(), outStream); - PaneInfoCoder.INSTANCE.encode(windowedElem.getPaneInfo(), outStream); + boolean metadataSupported = isMetadataSupported(); + PaneInfoCoder.INSTANCE.encode( + windowedElem.getPaneInfo().withElementMetadata(metadataSupported), outStream); + if (metadataSupported) { + BeamFnApi.Elements.ElementMetadata.Builder builder = + BeamFnApi.Elements.ElementMetadata.newBuilder(); + BeamFnApi.Elements.ElementMetadata em = + builder + .setDrain( + windowedElem.isDraining() != null + ? (Boolean.TRUE.equals(windowedElem.isDraining()) + ? BeamFnApi.Elements.DrainMode.Enum.DRAINING + : BeamFnApi.Elements.DrainMode.Enum.NOT_DRAINING) + : BeamFnApi.Elements.DrainMode.Enum.UNSPECIFIED) + .build(); + + ByteArrayCoder.of().encode(em.toByteArray(), outStream); + } valueCoder.encode(windowedElem.getValue(), outStream, context); } @@ -844,11 +936,21 @@ public WindowedValue decode(InputStream inStream, Context context) Instant timestamp = InstantCoder.of().decode(inStream); Collection windows = windowsCoder.decode(inStream); PaneInfo paneInfo = PaneInfoCoder.INSTANCE.decode(inStream); + Boolean draining = null; + if (isMetadataSupported() && paneInfo.isElementMetadata()) { + BeamFnApi.Elements.ElementMetadata elementMetadata = + BeamFnApi.Elements.ElementMetadata.parseFrom(ByteArrayCoder.of().decode(inStream)); + boolean b = elementMetadata.hasDrain(); + draining = + b + ? elementMetadata.getDrain().equals(BeamFnApi.Elements.DrainMode.Enum.DRAINING) + : null; + } T value = valueCoder.decode(inStream, context); // Because there are some remaining (incorrect) uses of WindowedValue with no windows, // we call this deprecated no-validation path when decoding - return WindowedValues.createWithoutValidation(value, timestamp, windows, paneInfo); + return WindowedValues.createWithoutValidation(value, timestamp, windows, paneInfo, draining); } @Override diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/PaneInfoTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/PaneInfoTest.java index 946deba036db..cda8ee1ea55c 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/PaneInfoTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/PaneInfoTest.java @@ -52,6 +52,22 @@ public void testEncodingRoundTrip() throws Exception { } } + @Test + public void testEncodingRoundTripWithElementMetadata() throws Exception { + Coder coder = PaneInfo.PaneInfoCoder.INSTANCE; + for (Timing timing : Timing.values()) { + long onTimeIndex = timing == Timing.EARLY ? -1 : 37; + CoderProperties.coderDecodeEncodeEqual( + coder, PaneInfo.createPane(false, false, timing, 389, onTimeIndex, true)); + CoderProperties.coderDecodeEncodeEqual( + coder, PaneInfo.createPane(false, true, timing, 5077, onTimeIndex, true)); + CoderProperties.coderDecodeEncodeEqual( + coder, PaneInfo.createPane(true, false, timing, 0, 0, true)); + CoderProperties.coderDecodeEncodeEqual( + coder, PaneInfo.createPane(true, true, timing, 0, 0, true)); + } + } + @Test public void testEncodings() { assertEquals( @@ -82,5 +98,9 @@ public void testEncodings() { "PaneInfo encoding should remain the same.", 0xF, PaneInfo.createPane(true, true, Timing.UNKNOWN).getEncodedByte()); + assertEquals( + "PaneInfo encoding should remain the same.", + 0x1, + PaneInfo.createPane(true, false, Timing.EARLY, 1, -1, true).getEncodedByte()); } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/WindowedValueTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/WindowedValueTest.java index 18660c5e6c36..3e3973e3720b 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/WindowedValueTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/WindowedValueTest.java @@ -77,6 +77,32 @@ public void testWindowedValueCoder() throws CoderException { Assert.assertArrayEquals(value.getWindows().toArray(), decodedValue.getWindows().toArray()); } + @Test + public void testWindowedValueWithElementMetadataCoder() throws CoderException { + WindowedValues.WindowedValueCoder.setMetadataSupported(); + Instant timestamp = new Instant(1234); + WindowedValue value = + WindowedValues.of( + "abc", + new Instant(1234), + Arrays.asList( + new IntervalWindow(timestamp, timestamp.plus(Duration.millis(1000))), + new IntervalWindow( + timestamp.plus(Duration.millis(1000)), timestamp.plus(Duration.millis(2000)))), + PaneInfo.NO_FIRING); + + Coder> windowedValueCoder = + WindowedValues.getFullCoder(StringUtf8Coder.of(), IntervalWindow.getCoder()); + + byte[] encodedValue = CoderUtils.encodeToByteArray(windowedValueCoder, value); + WindowedValue decodedValue = + CoderUtils.decodeFromByteArray(windowedValueCoder, encodedValue); + + Assert.assertEquals(value.getValue(), decodedValue.getValue()); + Assert.assertEquals(value.getTimestamp(), decodedValue.getTimestamp()); + Assert.assertArrayEquals(value.getWindows().toArray(), decodedValue.getWindows().toArray()); + } + @Test public void testFullWindowedValueCoderIsSerializableWithWellKnownCoderType() { CoderProperties.coderSerializable( diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/TimerTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/TimerTest.java index e96bf27c6ed9..a5d52a2dc892 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/TimerTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/TimerTest.java @@ -57,7 +57,8 @@ public void testTimer() { Collections.singleton(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING); + PaneInfo.NO_FIRING, + null); assertEquals("key", timer.getUserKey()); assertEquals("tag", timer.getDynamicTimerTag()); assertEquals(FIRE_TIME, timer.getFireTimestamp()); @@ -79,7 +80,8 @@ public void testTimerCoderWithInconsistentWithEqualsComponentCoders() throws Exc Collections.singleton(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING)); + PaneInfo.NO_FIRING, + null)); CoderProperties.structuralValueDecodeEncodeEqual( coder, Timer.cleared("key", "tag", Collections.singleton(GlobalWindow.INSTANCE))); CoderProperties.structuralValueConsistentWithEquals( @@ -90,14 +92,16 @@ public void testTimerCoderWithInconsistentWithEqualsComponentCoders() throws Exc Collections.singleton(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING), + PaneInfo.NO_FIRING, + null), Timer.of( "key", "tag", Collections.singleton(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING)); + PaneInfo.NO_FIRING, + null)); CoderProperties.structuralValueConsistentWithEquals( coder, Timer.cleared("key", "tag", Collections.singleton(GlobalWindow.INSTANCE)), @@ -115,7 +119,8 @@ public void testTimerCoderWithConsistentWithEqualsComponentCoders() throws Excep Collections.singletonList(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING)); + PaneInfo.NO_FIRING, + null)); CoderProperties.coderDecodeEncodeEqual( coder, Timer.cleared("key", "tag", Collections.singletonList(GlobalWindow.INSTANCE))); CoderProperties.coderConsistentWithEquals( @@ -126,14 +131,16 @@ public void testTimerCoderWithConsistentWithEqualsComponentCoders() throws Excep Collections.singletonList(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING), + PaneInfo.NO_FIRING, + null), Timer.of( "key", "tag", Collections.singletonList(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING)); + PaneInfo.NO_FIRING, + null)); CoderProperties.coderConsistentWithEquals( coder, Timer.cleared("key", "tag", Collections.singletonList(GlobalWindow.INSTANCE)), @@ -146,14 +153,16 @@ public void testTimerCoderWithConsistentWithEqualsComponentCoders() throws Excep Collections.singletonList(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING), + PaneInfo.NO_FIRING, + null), Timer.of( "key", "tag", Collections.singletonList(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING)); + PaneInfo.NO_FIRING, + null)); CoderProperties.coderDeterministic( coder, Timer.cleared("key", "tag", Collections.singletonList(GlobalWindow.INSTANCE)), diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java index 0388d3c03f00..4d32bcf53e27 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java @@ -247,6 +247,7 @@ public final void addRunnerForPTransform(Context context) throws IOException { private WindowedValue currentElement; private Object currentKey; + private Boolean draining; /** * Only valid during {@link @@ -1200,6 +1201,8 @@ private void processTimer( checkNotNull(timerBundleTracker); try { currentKey = timer.getUserKey(); + draining = timer.getDraining(); + // add drain Iterator windowIterator = (Iterator) timer.getWindows().iterator(); while (windowIterator.hasNext()) { @@ -1531,7 +1534,8 @@ private Timer getTimerForTime(Instant scheduledTime) { Collections.singletonList(boundedWindow), scheduledTime, outputTimestamp, - paneInfo); + paneInfo, + draining); } } @@ -1684,7 +1688,8 @@ public void output( Collections.singletonList(window), PaneInfo.NO_FIRING, currentRecordId, - currentRecordOffset)); + currentRecordOffset, + currentElement != null ? currentElement.isDraining() : null)); } @Override @@ -1708,7 +1713,8 @@ public void output( Collections.singletonList(window), PaneInfo.NO_FIRING, currentRecordId, - currentRecordOffset)); + currentRecordOffset, + currentElement != null ? currentElement.isDraining() : null)); } } @@ -1821,7 +1827,7 @@ public void outputWindowedValue( outputTo( mainOutputConsumer, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset, null)); } @Override @@ -1872,7 +1878,7 @@ public void outputWindowedValue( outputTo( consumer, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset, null)); } @Override @@ -1990,7 +1996,13 @@ public void outputWindowedValue( outputTo( mainOutputConsumer, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + currentElement != null ? currentElement.isDraining() : null)); } @Override @@ -2041,7 +2053,13 @@ public void outputWindowedValue( outputTo( consumer, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + currentElement != null ? currentElement.isDraining() : null)); } } @@ -2333,6 +2351,11 @@ public PaneInfo pane() { return currentElement.getPaneInfo(); } + @Override + public Boolean draining() { + return draining; + } + @Override public Object watermarkEstimatorState() { return currentWatermarkEstimatorState; @@ -2392,6 +2415,7 @@ public OutputBuilder builder(OutputT value) { .setWindow(currentWindow) .setTimestamp(currentTimer.getHoldTimestamp()) .setPaneInfo(currentTimer.getPaneInfo()) + .setDraining(draining) .setReceiver( windowedValue -> { checkOnWindowExpirationTimestamp(windowedValue.getTimestamp()); @@ -2411,7 +2435,13 @@ public void outputWindowedValue( outputTo( mainOutputConsumer, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + currentElement != null ? currentElement.isDraining() : null)); } @Override @@ -2468,7 +2498,13 @@ public void outputWindowedValue( outputTo( consumer, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + currentElement != null ? currentElement.isDraining() : null)); } @SuppressWarnings( @@ -2539,6 +2575,7 @@ public OutputBuilder builder(Row value) { .setValue(value) .setTimestamp(currentTimer.getHoldTimestamp()) .setWindow(currentWindow) + .setDraining(draining) .setReceiver( windowedValue -> context.outputWindowedValue( @@ -2576,6 +2613,7 @@ public OutputBuilder builder(T value) { .setValue(value) .setTimestamp(currentTimer.getHoldTimestamp()) .setWindow(currentWindow) + .setDraining(draining) .setReceiver( windowedValue -> context.outputWindowedValue( @@ -2613,6 +2651,7 @@ public OutputBuilder builder(Row value) { .setValue(value) .setTimestamp(currentTimer.getHoldTimestamp()) .setWindow(currentWindow) + .setDraining(draining) .setReceiver( windowedValue -> context.outputWindowedValue( @@ -2690,6 +2729,11 @@ public BoundedWindow window() { return currentWindow; } + @Override + public Boolean draining() { + return draining; + } + @Override public OutputBuilder builder(OutputT value) { return WindowedValues.builder() @@ -2697,6 +2741,7 @@ public OutputBuilder builder(OutputT value) { .setTimestamp(currentTimer.getHoldTimestamp()) .setWindow(currentWindow) .setPaneInfo(currentTimer.getPaneInfo()) + .setDraining(draining) .setReceiver( windowedValue -> { checkTimerTimestamp(windowedValue.getTimestamp()); @@ -2735,7 +2780,13 @@ public void outputWindowedValue( outputTo( mainOutputConsumer, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + currentElement != null ? currentElement.isDraining() : null)); } @Override @@ -2903,6 +2954,7 @@ public OutputBuilder builder(T value) { .setValue(value) .setTimestamp(currentTimer.getHoldTimestamp()) .setWindow(currentWindow) + .setDraining(draining) .setPaneInfo(currentTimer.getPaneInfo()) .setReceiver( windowedValue -> @@ -2941,6 +2993,7 @@ public OutputBuilder builder(Row value) { .setTimestamp(currentTimer.getHoldTimestamp()) .setWindow(currentWindow) .setPaneInfo(currentTimer.getPaneInfo()) + .setDraining(draining) .setReceiver( windowedValue -> context.outputWindowedValue( diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java index ef19b7c18804..c68dc0ba8eac 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java @@ -1161,7 +1161,8 @@ private org.apache.beam.sdk.util.construction.Timer dynamicTimerInGlobalW Collections.singletonList(GlobalWindow.INSTANCE), fireTimestamp, holdTimestamp, - PaneInfo.NO_FIRING); + PaneInfo.NO_FIRING, + null); } private WindowedValue valueInWindows( diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/SplittablePairWithRestrictionDoFnRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/SplittablePairWithRestrictionDoFnRunnerTest.java index 13a0b105ec39..8101d6dcc83a 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/SplittablePairWithRestrictionDoFnRunnerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/SplittablePairWithRestrictionDoFnRunnerTest.java @@ -123,7 +123,8 @@ private org.apache.beam.sdk.util.construction.Timer dynamicTimerInGlobalW Collections.singletonList(GlobalWindow.INSTANCE), fireTimestamp, holdTimestamp, - PaneInfo.NO_FIRING); + PaneInfo.NO_FIRING, + null); } private WindowedValue valueInWindows( diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index a7a62571e38e..11823886759c 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -1112,7 +1112,8 @@ public void testInstructionEmbeddedElementsAreProcessed() throws Exception { Collections.singletonList(GlobalWindow.INSTANCE), Instant.ofEpochMilli(1L), Instant.ofEpochMilli(1L), - PaneInfo.ON_TIME_AND_ONLY_FIRING), + PaneInfo.ON_TIME_AND_ONLY_FIRING, + null), encodedTimer); Elements elements = Elements.newBuilder() @@ -1231,7 +1232,8 @@ public void testInstructionEmbeddedElementsWithMalformedTimers() throws Exceptio Collections.singletonList(GlobalWindow.INSTANCE), Instant.ofEpochMilli(1L), Instant.ofEpochMilli(1L), - PaneInfo.ON_TIME_AND_ONLY_FIRING), + PaneInfo.ON_TIME_AND_ONLY_FIRING, + null), encodedTimer); assertThrows( @@ -1325,7 +1327,8 @@ public void testOutputEmbeddedElementsAreProcessed() throws Exception { Collections.singletonList(GlobalWindow.INSTANCE), Instant.ofEpochMilli(1L), Instant.ofEpochMilli(1L), - PaneInfo.ON_TIME_AND_ONLY_FIRING), + PaneInfo.ON_TIME_AND_ONLY_FIRING, + null), encodedTimer); InstructionResponse.Builder builder =