Skip to content

Commit 1c0c024

Browse files
acritesacrites
authored
Adds backlog reporting support for non-fnapi based SDF's. (#38346)
1 parent 60620e7 commit 1c0c024

11 files changed

Lines changed: 415 additions & 128 deletions

File tree

‎.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming_Engine.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
{
22
"comment": "Modify this file in a trivial way to cause this test suite to run!",
3-
"modification": 2,
3+
"modification": 1,
44
}

‎runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java‎

Lines changed: 138 additions & 109 deletions
Original file line numberDiff line numberDiff line change
@@ -130,114 +130,114 @@ public Result invokeProcessElement(
130130
final Map<String, PCollectionView<?>> sideInputMapping) {
131131
final ProcessContext processContext = new ProcessContext(element, tracker, watermarkEstimator);
132132

133-
DoFn.ProcessContinuation cont =
134-
invoker.invokeProcessElement(
135-
new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
136-
@Override
137-
public String getErrorContext() {
138-
return OutputAndTimeBoundedSplittableProcessElementInvoker.class.getSimpleName();
139-
}
140-
141-
@Override
142-
public DoFn<InputT, OutputT>.ProcessContext processContext(
143-
DoFn<InputT, OutputT> doFn) {
144-
return processContext;
145-
}
146-
147-
@Override
148-
public Object sideInput(String tagId) {
149-
PCollectionView<?> view = sideInputMapping.get(tagId);
150-
if (view == null) {
151-
throw new IllegalArgumentException("calling getSideInput() with unknown view");
152-
}
153-
return processContext.sideInput(view);
154-
}
155-
156-
@Override
157-
public Object restriction() {
158-
return tracker.currentRestriction();
159-
}
160-
161-
@Override
162-
public InputT element(DoFn<InputT, OutputT> doFn) {
163-
return processContext.element();
164-
}
165-
166-
@Override
167-
public Instant timestamp(DoFn<InputT, OutputT> doFn) {
168-
return processContext.timestamp();
169-
}
170-
171-
@Override
172-
public String timerId(DoFn<InputT, OutputT> doFn) {
173-
throw new UnsupportedOperationException(
174-
"Cannot access timerId as parameter outside of @OnTimer method.");
175-
}
176-
177-
@Override
178-
public TimeDomain timeDomain(DoFn<InputT, OutputT> doFn) {
179-
throw new UnsupportedOperationException(
180-
"Access to time domain not supported in ProcessElement");
181-
}
182-
183-
@Override
184-
public OutputReceiver<OutputT> outputReceiver(DoFn<InputT, OutputT> doFn) {
185-
return DoFnOutputReceivers.windowedReceiver(
186-
processContext, OutputBuilderSuppliers.supplierForElement(element), null);
187-
}
188-
189-
@Override
190-
public OutputReceiver<Row> outputRowReceiver(DoFn<InputT, OutputT> doFn) {
191-
throw new UnsupportedOperationException("Not supported in SplittableDoFn");
192-
}
193-
194-
@Override
195-
public MultiOutputReceiver taggedOutputReceiver(DoFn<InputT, OutputT> doFn) {
196-
return DoFnOutputReceivers.windowedMultiReceiver(
197-
processContext, OutputBuilderSuppliers.supplierForElement(element));
198-
}
199-
200-
@Override
201-
public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
202-
return processContext.causedByDrain();
203-
}
204-
205-
@Override
206-
public RestrictionTracker<?, ?> restrictionTracker() {
207-
return processContext.tracker;
208-
}
209-
210-
@Override
211-
public WatermarkEstimator<?> watermarkEstimator() {
212-
return processContext.watermarkEstimator;
213-
}
214-
215-
@Override
216-
public PipelineOptions pipelineOptions() {
217-
return pipelineOptions;
218-
}
219-
220-
@Override
221-
public BundleFinalizer bundleFinalizer() {
222-
return bundleFinalizer.get();
223-
}
224-
225-
// Unsupported methods below.
226-
227-
@Override
228-
public StartBundleContext startBundleContext(DoFn<InputT, OutputT> doFn) {
229-
throw new IllegalStateException(
230-
"Should not access startBundleContext() from @"
231-
+ DoFn.ProcessElement.class.getSimpleName());
232-
}
233-
234-
@Override
235-
public FinishBundleContext finishBundleContext(DoFn<InputT, OutputT> doFn) {
236-
throw new IllegalStateException(
237-
"Should not access finishBundleContext() from @"
238-
+ DoFn.ProcessElement.class.getSimpleName());
239-
}
240-
});
133+
DoFnInvoker.BaseArgumentProvider<InputT, OutputT> invokerArgumentProvider =
134+
new DoFnInvoker.BaseArgumentProvider<InputT, OutputT>() {
135+
@Override
136+
public String getErrorContext() {
137+
return OutputAndTimeBoundedSplittableProcessElementInvoker.class.getSimpleName();
138+
}
139+
140+
@Override
141+
public DoFn<InputT, OutputT>.ProcessContext processContext(DoFn<InputT, OutputT> doFn) {
142+
return processContext;
143+
}
144+
145+
@Override
146+
public Object sideInput(String tagId) {
147+
PCollectionView<?> view = sideInputMapping.get(tagId);
148+
if (view == null) {
149+
throw new IllegalArgumentException("calling getSideInput() with unknown view");
150+
}
151+
return processContext.sideInput(view);
152+
}
153+
154+
@Override
155+
public Object restriction() {
156+
return tracker.currentRestriction();
157+
}
158+
159+
@Override
160+
public InputT element(DoFn<InputT, OutputT> doFn) {
161+
return processContext.element();
162+
}
163+
164+
@Override
165+
public Instant timestamp(DoFn<InputT, OutputT> doFn) {
166+
return processContext.timestamp();
167+
}
168+
169+
@Override
170+
public String timerId(DoFn<InputT, OutputT> doFn) {
171+
throw new UnsupportedOperationException(
172+
"Cannot access timerId as parameter outside of @OnTimer method.");
173+
}
174+
175+
@Override
176+
public TimeDomain timeDomain(DoFn<InputT, OutputT> doFn) {
177+
throw new UnsupportedOperationException(
178+
"Access to time domain not supported in ProcessElement");
179+
}
180+
181+
@Override
182+
public OutputReceiver<OutputT> outputReceiver(DoFn<InputT, OutputT> doFn) {
183+
return DoFnOutputReceivers.windowedReceiver(
184+
processContext, OutputBuilderSuppliers.supplierForElement(element), null);
185+
}
186+
187+
@Override
188+
public OutputReceiver<Row> outputRowReceiver(DoFn<InputT, OutputT> doFn) {
189+
throw new UnsupportedOperationException("Not supported in SplittableDoFn");
190+
}
191+
192+
@Override
193+
public MultiOutputReceiver taggedOutputReceiver(DoFn<InputT, OutputT> doFn) {
194+
return DoFnOutputReceivers.windowedMultiReceiver(
195+
processContext, OutputBuilderSuppliers.supplierForElement(element));
196+
}
197+
198+
@Override
199+
public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
200+
return processContext.causedByDrain();
201+
}
202+
203+
@Override
204+
public RestrictionTracker<?, ?> restrictionTracker() {
205+
return processContext.tracker;
206+
}
207+
208+
@Override
209+
public WatermarkEstimator<?> watermarkEstimator() {
210+
return processContext.watermarkEstimator;
211+
}
212+
213+
@Override
214+
public PipelineOptions pipelineOptions() {
215+
return pipelineOptions;
216+
}
217+
218+
@Override
219+
public BundleFinalizer bundleFinalizer() {
220+
return bundleFinalizer.get();
221+
}
222+
223+
// Unsupported methods below.
224+
225+
@Override
226+
public StartBundleContext startBundleContext(DoFn<InputT, OutputT> doFn) {
227+
throw new IllegalStateException(
228+
"Should not access startBundleContext() from @"
229+
+ DoFn.ProcessElement.class.getSimpleName());
230+
}
231+
232+
@Override
233+
public FinishBundleContext finishBundleContext(DoFn<InputT, OutputT> doFn) {
234+
throw new IllegalStateException(
235+
"Should not access finishBundleContext() from @"
236+
+ DoFn.ProcessElement.class.getSimpleName());
237+
}
238+
};
239+
240+
DoFn.ProcessContinuation cont = invoker.invokeProcessElement(invokerArgumentProvider);
241241
processContext.cancelScheduledCheckpoint();
242242
@Nullable
243243
KV<RestrictionT, KV<Instant, WatermarkEstimatorStateT>> residual =
@@ -278,8 +278,37 @@ public FinishBundleContext finishBundleContext(DoFn<InputT, OutputT> doFn) {
278278
if (residual == null) {
279279
return new Result(null, cont, null, null);
280280
}
281+
final KV<RestrictionT, KV<Instant, WatermarkEstimatorStateT>> residualForGetSize = residual;
282+
// For a list of all DoFnInvoker arguments, see DoFn.java.
283+
double backlogBytes =
284+
invoker.invokeGetSize(
285+
new DoFnInvoker.DelegatingArgumentProvider<InputT, OutputT>(
286+
invokerArgumentProvider, invokerArgumentProvider.getErrorContext() + "/GetSize") {
287+
@Override
288+
public Object restriction() {
289+
return residualForGetSize.getKey();
290+
}
291+
292+
@Override
293+
public RestrictionTracker<?, ?> restrictionTracker() {
294+
return invoker.invokeNewTracker(
295+
new DoFnInvoker.DelegatingArgumentProvider<InputT, OutputT>(
296+
invokerArgumentProvider,
297+
invokerArgumentProvider.getErrorContext() + "/NewTracker") {
298+
299+
@Override
300+
public Object restriction() {
301+
return residualForGetSize.getKey();
302+
}
303+
});
304+
}
305+
});
281306
return new Result(
282-
residual.getKey(), cont, residual.getValue().getKey(), residual.getValue().getValue());
307+
residual.getKey(),
308+
cont,
309+
residual.getValue().getKey(),
310+
residual.getValue().getValue(),
311+
backlogBytes);
283312
}
284313

285314
private class ProcessContext extends DoFn<InputT, OutputT>.ProcessContext

‎runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import com.google.auto.service.AutoService;
2323
import java.util.List;
2424
import java.util.Map;
25+
import java.util.function.Consumer;
2526
import org.apache.beam.model.pipeline.v1.RunnerApi;
2627
import org.apache.beam.sdk.coders.ByteArrayCoder;
2728
import org.apache.beam.sdk.coders.Coder;
@@ -281,6 +282,7 @@ public static class ProcessFn<InputT, OutputT, RestrictionT, PositionT, Watermar
281282
processElementInvoker;
282283

283284
private transient @Nullable DoFnInvoker<InputT, OutputT> invoker;
285+
private transient @Nullable Consumer<Double> backlogBytesCallback;
284286

285287
public ProcessFn(
286288
DoFn<InputT, OutputT> fn,
@@ -323,6 +325,10 @@ public void setProcessElementInvoker(
323325
this.processElementInvoker = invoker;
324326
}
325327

328+
public void setBacklogBytesCallback(Consumer<Double> backlogBytesCallback) {
329+
this.backlogBytesCallback = backlogBytesCallback;
330+
}
331+
326332
public DoFn<InputT, OutputT> getFn() {
327333
return fn;
328334
}
@@ -624,6 +630,9 @@ public String getErrorContext() {
624630
} else {
625631
holdState.clear();
626632
}
633+
if (backlogBytesCallback != null && result.getBacklogBytes() >= 0) {
634+
backlogBytesCallback.accept(result.getBacklogBytes());
635+
}
627636
}
628637

629638
private DoFnInvoker.ArgumentProvider<InputT, OutputT> wrapOptionsAsSetup(

‎runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableProcessElementInvoker.java‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,10 @@ public class Result {
4242
private final DoFn.ProcessContinuation continuation;
4343
private final @Nullable Instant futureOutputWatermark;
4444
private final @Nullable WatermarkEstimatorStateT futureWatermarkEstimatorState;
45+
private final double backlogBytes;
46+
47+
/* Constant representing an unknown amount of backlog. */
48+
public static final double BACKLOG_UNKNOWN = -1.0;
4549

4650
@SuppressFBWarnings(
4751
value = "NP_PARAMETER_MUST_BE_NONNULL_BUT_MARKED_AS_NULLABLE",
@@ -50,12 +54,27 @@ public Result(
5054
@Nullable RestrictionT residualRestriction,
5155
DoFn.ProcessContinuation continuation,
5256
@Nullable Instant futureOutputWatermark,
53-
@Nullable WatermarkEstimatorStateT futureWatermarkEstimatorState) {
57+
@Nullable WatermarkEstimatorStateT futureWatermarkEstimatorState,
58+
double backlogBytes) {
5459
checkArgument(continuation != null, "continuation must not be null");
5560
this.continuation = continuation;
5661
this.residualRestriction = residualRestriction;
5762
this.futureOutputWatermark = futureOutputWatermark;
5863
this.futureWatermarkEstimatorState = futureWatermarkEstimatorState;
64+
this.backlogBytes = backlogBytes;
65+
}
66+
67+
public Result(
68+
@Nullable RestrictionT residualRestriction,
69+
DoFn.ProcessContinuation continuation,
70+
@Nullable Instant futureOutputWatermark,
71+
@Nullable WatermarkEstimatorStateT futureWatermarkEstimatorState) {
72+
this(
73+
residualRestriction,
74+
continuation,
75+
futureOutputWatermark,
76+
futureWatermarkEstimatorState,
77+
BACKLOG_UNKNOWN);
5978
}
6079

6180
/**
@@ -76,6 +95,10 @@ public DoFn.ProcessContinuation getContinuation() {
7695
public @Nullable WatermarkEstimatorStateT getFutureWatermarkEstimatorState() {
7796
return futureWatermarkEstimatorState;
7897
}
98+
99+
public double getBacklogBytes() {
100+
return backlogBytes;
101+
}
79102
}
80103

81104
/**

‎runners/core-java/src/main/java/org/apache/beam/runners/core/StepContext.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,4 +36,10 @@ public interface StepContext {
3636
default BundleFinalizer bundleFinalizer() {
3737
throw new UnsupportedOperationException("BundleFinalizer is unsupported.");
3838
}
39+
40+
/**
41+
* Set the current backlog bytes for this step. This is mainly used by splittable DoFn to report
42+
* the size of the residual restriction.
43+
*/
44+
default void setBacklogBytes(double backlogBytes) {}
3945
}

0 commit comments

Comments
 (0)