[Drain] Expose drain to dofn processElement and onTimer - add missing implementation for WindowObservingProcessBundleContextBase and SDF

Author: stankiewicz

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java

@@ -1718,6 +1718,11 @@ public BoundedWindow window() {
       return currentWindow;
     }
 
+    @Override
+    public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
+      return currentElement.causedByDrain();
+    }
+
     @Override
     public Object sideInput(String tagId) {
       return sideInput(sideInputMapping.get(tagId));
@@ -2422,6 +2427,11 @@ public BoundedWindow window() {
       return currentWindow;
     }
 
+    @Override
+    public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn){
+      return currentTimer.causedByDrain();
+    }
+
     @Override
     public Instant timestamp(DoFn<InputT, OutputT> doFn) {
       return currentTimer.getHoldTimestamp();

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplittablePairWithRestrictionDoFnRunner.java

@@ -18,6 +18,7 @@
 package org.apache.beam.fn.harness;
 
 import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+import org.apache.beam.sdk.values.CausedByDrain;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
 
 import com.google.auto.service.AutoService;
@@ -278,6 +279,11 @@ public PaneInfo paneInfo(DoFn<InputT, OutputT> doFn) {
       return getCurrentElementOrFail().getPaneInfo();
     }
 
+    @Override
+    public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
+      return getCurrentElementOrFail().causedByDrain();
+    }
+
     @Override
     public Object schemaElement(int index) {
       SerializableFunction<InputT, Object> converter =