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

Author: stankiewicz

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

@@ -2422,6 +2422,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

@@ -41,6 +41,7 @@
 import org.apache.beam.sdk.util.UserCodeException;
 import org.apache.beam.sdk.util.construction.PTransformTranslation;
 import org.apache.beam.sdk.util.construction.ParDoTranslation;
+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;
@@ -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 =