Skip to content

[Drain] OnTimer - propagate caused by drain bit up to DoFnRunner #37012

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
stankiewicz wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[Drain] OnTimer - propagate caused by drain bit up to DoFnRunner #37012

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.CausedByDrain;
import org.apache.beam.sdk.values.WindowedValue;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Instant;
Expand All @@ -46,7 +47,8 @@ public interface DoFnRunner<InputT extends @Nullable Object, OutputT extends @Nu
BoundedWindow window,
Instant timestamp,
Instant outputTimestamp,
TimeDomain timeDomain);
TimeDomain timeDomain,
CausedByDrain causedByDrain);

/**
* Calls a {@link DoFn DoFn's} {@link DoFn.FinishBundle @FinishBundle} method and performs
Expand Down
7 changes: 5 additions & 2 deletions runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.WindowTracing;
import org.apache.beam.sdk.values.CausedByDrain;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.sdk.values.WindowedValues;
Expand Down Expand Up @@ -89,8 +90,10 @@ public <KeyT> void onTimer(
BoundedWindow window,
Instant timestamp,
Instant outputTimestamp,
TimeDomain timeDomain) {
doFnRunner.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain);
TimeDomain timeDomain,
CausedByDrain causedByDrain) {
doFnRunner.onTimer(
timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, causedByDrain);
}

@Override
Expand Down
6 changes: 6 additions & 0 deletions ...ava/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.util.OutputBuilderSuppliers;
import org.apache.beam.sdk.util.WindowedValueMultiReceiver;
import org.apache.beam.sdk.values.CausedByDrain;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.Row;
Expand Down Expand Up @@ -396,6 +397,11 @@ public PaneInfo pane() {
return element.getRecordOffset();
}

@Override
public CausedByDrain causedByDrain() {
return CausedByDrain.NORMAL;
Copy link
Member

@kennknowles kennknowles Feb 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this come from the element metadata?

Copy link
Member

@kennknowles kennknowles Feb 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this come from the element metadata?

}

@Override
public PipelineOptions getPipelineOptions() {
return pipelineOptions;
Expand Down
25 changes: 21 additions & 4 deletions runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@
import org.apache.beam.sdk.util.SystemDoFnInternal;
import org.apache.beam.sdk.util.UserCodeException;
import org.apache.beam.sdk.util.WindowedValueMultiReceiver;
import org.apache.beam.sdk.values.CausedByDrain;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.sdk.values.TupleTag;
Expand Down Expand Up @@ -200,11 +201,13 @@ public <KeyT> void onTimer(
BoundedWindow window,
Instant timestamp,
Instant outputTimestamp,
TimeDomain timeDomain) {
TimeDomain timeDomain,
CausedByDrain causedByDrain) {
Preconditions.checkNotNull(outputTimestamp, "outputTimestamp");

OnTimerArgumentProvider<KeyT> argumentProvider =
new OnTimerArgumentProvider<>(timerId, key, window, timestamp, outputTimestamp, timeDomain);
new OnTimerArgumentProvider<>(
timerId, key, window, timestamp, outputTimestamp, timeDomain, causedByDrain);
invoker.invokeOnTimer(timerId, timerFamilyId, argumentProvider);
}

Expand Down Expand Up @@ -399,6 +402,11 @@ public InputT element() {
return elem.getValue();
}

@Override
public CausedByDrain causedByDrain() {
return elem.causedByDrain();
}

@Override
public <T> T sideInput(PCollectionView<T> view) {
checkNotNull(view, "View passed to sideInput cannot be null");
Expand Down Expand Up @@ -702,6 +710,7 @@ private class OnTimerArgumentProvider<KeyT> extends DoFn<InputT, OutputT>.OnTime
private final TimeDomain timeDomain;
private final String timerId;
private final KeyT key;
private final CausedByDrain causedByDrain;
private final OutputBuilderSupplier builderSupplier;

/** Lazily initialized; should only be accessed via {@link #getNamespace()}. */
Expand All @@ -727,28 +736,36 @@ private OnTimerArgumentProvider(
BoundedWindow window,
Instant fireTimestamp,
Instant timestamp,
TimeDomain timeDomain) {
TimeDomain timeDomain,
CausedByDrain causedByDrain) {
fn.super();
this.timerId = timerId;
this.window = window;
this.fireTimestamp = fireTimestamp;
this.timestamp = timestamp;
this.timeDomain = timeDomain;
this.key = key;
this.causedByDrain = causedByDrain;
this.builderSupplier =
OutputBuilderSuppliers.supplierForElement(
WindowedValues.builder()
.setValue(null)
.setTimestamp(timestamp)
.setWindow(window)
.setPaneInfo(PaneInfo.NO_FIRING));
.setPaneInfo(PaneInfo.NO_FIRING)
.setCausedByDrain(causedByDrain));
}

@Override
public Instant timestamp() {
return timestamp;
}

@Override
public CausedByDrain causedByDrain() {
return causedByDrain;
}

@Override
public Instant fireTimestamp() {
return fireTimestamp;
Expand Down
11 changes: 10 additions & 1 deletion ...re-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.CausedByDrain;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
Expand Down Expand Up @@ -116,7 +117,15 @@ public <KeyT> 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,
CausedByDrain.NORMAL);
Copy link
Member

@kennknowles kennknowles Feb 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it should propagate here? will that be a follow-up that adds it to the pushback side input DoFnRunner?

}

@Override
Expand Down
16 changes: 13 additions & 3 deletions runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
import org.apache.beam.sdk.transforms.windowing.NonMergingWindowFn;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.util.WindowTracing;
import org.apache.beam.sdk.values.CausedByDrain;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.sdk.values.WindowedValues;
import org.apache.beam.sdk.values.WindowingStrategy;
Expand All @@ -48,7 +49,8 @@
/**
* A customized {@link DoFnRunner} that handles late data dropping and garbage collection for
* stateful {@link DoFn DoFns}. It registers a GC timer in {@link #processElement(WindowedValue)}
* and does cleanup in {@link #onTimer(String, String, BoundedWindow, Instant, Instant, TimeDomain)}
* and does cleanup in {@link #onTimer(String, String, BoundedWindow, Instant, Instant, TimeDomain,
* boolean)}
*
* @param <InputT> the type of the {@link DoFn} (main) input elements
* @param <OutputT> the type of the {@link DoFn} (main) output elements
Expand Down Expand Up @@ -208,7 +210,8 @@ public <KeyT> void onTimer(
BoundedWindow window,
Instant timestamp,
Instant outputTimestamp,
TimeDomain timeDomain) {
TimeDomain timeDomain,
CausedByDrain causedByDrain) {

if (timerId.equals(SORT_FLUSH_TIMER)) {
onSortFlushTimer(window, stepContext.timerInternals().currentInputWatermarkTime());
Expand All @@ -232,7 +235,14 @@ public <KeyT> void onTimer(
stepContext.timerInternals().currentInputWatermarkTime());
} else {
doFnRunner.onTimer(
timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain);
timerId,
timerFamilyId,
key,
window,
timestamp,
outputTimestamp,
timeDomain,
causedByDrain);
}
}
}
Expand Down
17 changes: 11 additions & 6 deletions runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,8 @@ public void testOnTimerExceptionsWrappedAsUserCodeException() {
GlobalWindow.INSTANCE,
new Instant(0),
new Instant(0),
TimeDomain.EVENT_TIME);
TimeDomain.EVENT_TIME,
CausedByDrain.NORMAL);
Copy link
Member

@kennknowles kennknowles Feb 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe this unit test file should test caused by drain propagation somehow?

}

/**
Expand Down Expand Up @@ -266,7 +267,8 @@ public void testOnTimerCalled() {
GlobalWindow.INSTANCE,
currentTime.plus(offset),
currentTime.plus(offset),
TimeDomain.EVENT_TIME);
TimeDomain.EVENT_TIME,
CausedByDrain.NORMAL);

assertThat(
fn.onTimerInvocations,
Expand All @@ -277,7 +279,8 @@ public void testOnTimerCalled() {
StateNamespaces.window(windowFn.windowCoder(), GlobalWindow.INSTANCE),
currentTime.plus(offset),
currentTime.plus(offset),
TimeDomain.EVENT_TIME)));
TimeDomain.EVENT_TIME,
CausedByDrain.NORMAL)));
}

/**
Expand Down Expand Up @@ -593,7 +596,8 @@ public void testOnTimerAllowedSkew() {
GlobalWindow.INSTANCE,
new Instant(0),
new Instant(0),
TimeDomain.EVENT_TIME);
TimeDomain.EVENT_TIME,
CausedByDrain.NORMAL);
}

@Test
Expand Down Expand Up @@ -625,7 +629,8 @@ public void testOnTimerNoSkew() {
GlobalWindow.INSTANCE,
new Instant(0),
new Instant(0),
TimeDomain.EVENT_TIME);
TimeDomain.EVENT_TIME,
CausedByDrain.NORMAL);
});

assertThat(exception.getCause(), isA(IllegalArgumentException.class));
Expand Down Expand Up @@ -703,7 +708,7 @@ public void onTimer(OnTimerContext context) {
context.fireTimestamp(),
context.timestamp(),
context.timeDomain(),
CausedByDrain.NORMAL));
context.causedByDrain()));
}
}

Expand Down
3 changes: 2 additions & 1 deletion ...ava/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -353,7 +353,8 @@ public <KeyT> void onTimer(
BoundedWindow window,
Instant timestamp,
Instant outputTimestamp,
TimeDomain timeDomain) {
TimeDomain timeDomain,
CausedByDrain causedByDrain) {
firedTimers.add(
TimerData.of(
timerId,
Expand Down
3 changes: 2 additions & 1 deletion runners/core-java/src/test/java/org/apache/beam/runners/core/StatefulDoFnRunnerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -468,7 +468,8 @@ private static void advanceInputWatermark(
window,
timer.getTimestamp(),
timer.getOutputTimestamp(),
timer.getDomain());
timer.getDomain(),
timer.causedByDrain());
}
}

Expand Down
3 changes: 2 additions & 1 deletion ...n/java/org/apache/beam/runners/flink/translation/functions/FlinkStatefulDoFnFunction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,8 @@ private void fireTimer(
window,
timer.getTimestamp(),
timer.getOutputTimestamp(),
timer.getDomain());
timer.getDomain(),
timer.causedByDrain());
}

@Override
Expand Down
14 changes: 12 additions & 2 deletions ...link/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.CausedByDrain;
import org.apache.beam.sdk.values.WindowedValue;
import org.joda.time.Instant;

Expand Down Expand Up @@ -73,10 +74,19 @@ public <KeyT> void onTimer(
final BoundedWindow window,
final Instant timestamp,
final Instant outputTimestamp,
final TimeDomain timeDomain) {
final TimeDomain timeDomain,
CausedByDrain causedByDrain) {
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,
causedByDrain);
} catch (IOException e) {
throw new RuntimeException(e);
}
Expand Down
3 changes: 2 additions & 1 deletion ...n/java/org/apache/beam/runners/flink/translation/functions/FlinkStatefulDoFnFunction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,8 @@ private void fireTimer(
window,
timer.getTimestamp(),
timer.getOutputTimestamp(),
timer.getDomain());
timer.getDomain(),
timer.causedByDrain());
}

@Override
Expand Down
7 changes: 5 additions & 2 deletions ...apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,7 @@
import org.apache.beam.sdk.util.construction.Timer;
import org.apache.beam.sdk.util.construction.graph.ExecutableStage;
import org.apache.beam.sdk.util.construction.graph.UserStateReference;
import org.apache.beam.sdk.values.CausedByDrain;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
Expand Down Expand Up @@ -1001,7 +1002,8 @@ public <KeyT> void onTimer(
BoundedWindow window,
Instant timestamp,
Instant outputTimestamp,
TimeDomain timeDomain) {
TimeDomain timeDomain,
CausedByDrain causedByDrain) {
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");
Expand Down Expand Up @@ -1034,7 +1036,8 @@ public <KeyT> void onTimer(
timestamp,
outputTimestamp,
// TODO: Support propagating the PaneInfo through.
PaneInfo.NO_FIRING);
PaneInfo.NO_FIRING,
causedByDrain);
try {
timerReceiver.accept(timerValue);
} catch (Exception e) {
Expand Down
10 changes: 9 additions & 1 deletion ...pache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.CausedByDrain;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
import org.checkerframework.checker.nullness.qual.Nullable;
Expand Down Expand Up @@ -104,7 +105,14 @@ static final class Timer<KeyT> implements BufferedElement {
@Override
public void processWith(DoFnRunner doFnRunner) {
doFnRunner.onTimer(
timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain);
timerId,
timerFamilyId,
key,
window,
timestamp,
outputTimestamp,
timeDomain,
CausedByDrain.NORMAL);
}

@Override
Expand Down
Loading
Loading