Skip to content
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

[BEAM-13015] Only create a TimerBundleTracker if there are timers. #17445

Merged
merged 1 commit into from
May 20, 2022
Merged

[BEAM-13015] Only create a TimerBundleTracker if there are timers. #17445

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
36 changes: 24 additions & 12 deletions sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,7 @@ static class Factory<InputT, RestrictionT, PositionT, WatermarkEstimatorStateT,
private final String mainInputId;
private final FnApiStateAccessor<?> stateAccessor;
private final Map<String, FnDataReceiver<?>> outboundTimerReceivers;
private FnApiTimerBundleTracker timerBundleTracker;
private final @Nullable FnApiTimerBundleTracker timerBundleTracker;
private final DoFnInvoker<InputT, OutputT> doFnInvoker;
private final StartBundleArgumentProvider startBundleArgumentProvider;
private final ProcessBundleContextBase processContext;
Expand Down Expand Up @@ -483,10 +483,6 @@ private interface TriFunction<FirstT, SecondT, ThirdT> {
this.bundleFinalizer = bundleFinalizer;
this.onTimerContext = new OnTimerContext();
this.onWindowExpirationContext = new OnWindowExpirationContext<>();
this.timerBundleTracker =
new FnApiTimerBundleTracker(
keyCoder, windowCoder, this::getCurrentKey, () -> currentWindow);
addResetFunction.accept(timerBundleTracker::reset);

this.mainOutputConsumers =
(Collection<FnDataReceiver<WindowedValue<OutputT>>>)
Expand Down Expand Up @@ -737,12 +733,22 @@ private ByteString encodeProgress(double value) throws IOException {
() -> currentWindow);

// Register as a consumer for each timer.
outboundTimerReceivers = new HashMap<>();
for (Map.Entry<String, KV<TimeDomain, Coder<Timer<Object>>>> timerFamilyInfo :
timerFamilyInfos.entrySet()) {
String localName = timerFamilyInfo.getKey();
Coder<Timer<Object>> timerCoder = timerFamilyInfo.getValue().getValue();
outboundTimerReceivers.put(localName, getOutgoingTimersConsumer.apply(localName, timerCoder));
this.outboundTimerReceivers = new HashMap<>();
if (timerFamilyInfos.isEmpty()) {
this.timerBundleTracker = null;
} else {
this.timerBundleTracker =
new FnApiTimerBundleTracker(
keyCoder, windowCoder, this::getCurrentKey, () -> currentWindow);
addResetFunction.accept(timerBundleTracker::reset);

for (Map.Entry<String, KV<TimeDomain, Coder<Timer<Object>>>> timerFamilyInfo :
timerFamilyInfos.entrySet()) {
String localName = timerFamilyInfo.getKey();
Coder<Timer<Object>> timerCoder = timerFamilyInfo.getValue().getValue();
outboundTimerReceivers.put(
localName, getOutgoingTimersConsumer.apply(localName, timerCoder));
}
}
}

Expand Down Expand Up @@ -1639,6 +1645,7 @@ private HandlesSplits.SplitResult trySplitForElementAndRestriction(

private <K> void processTimer(
String timerIdOrTimerFamilyId, TimeDomain timeDomain, Timer<K> timer) {
checkNotNull(timerBundleTracker);
try {
currentKey = timer.getUserKey();
Iterator<BoundedWindow> windowIterator =
Expand Down Expand Up @@ -1737,7 +1744,9 @@ private <K> void processOnWindowExpiration(Timer<K> timer) {
}

private void finishBundle() throws Exception {
timerBundleTracker.outputTimers(outboundTimerReceivers::get);
if (timerBundleTracker != null) {
timerBundleTracker.outputTimers(outboundTimerReceivers::get);
}

doFnInvoker.invokeFinishBundle(finishBundleArgumentProvider);

Expand Down Expand Up @@ -1811,6 +1820,7 @@ private class FnApiTimer<K> implements org.apache.beam.sdk.state.Timer {

@Override
public void set(Instant absoluteTime) {
checkNotNull(timerBundleTracker);
// Ensures that the target time is reasonable. For event time timers this means that the time
// should be prior to window GC time.
if (TimeDomain.EVENT_TIME.equals(timeDomain)) {
Expand All @@ -1827,6 +1837,7 @@ public void set(Instant absoluteTime) {

@Override
public void setRelative() {
checkNotNull(timerBundleTracker);
Instant target;
if (period.equals(Duration.ZERO)) {
target = fireTimestamp.plus(offset);
Expand All @@ -1843,6 +1854,7 @@ public void setRelative() {

@Override
public void clear() {
checkNotNull(timerBundleTracker);
timerBundleTracker.timerModified(timerIdOrFamily, timeDomain, getClearedTimer());
}

Expand Down