ontimer

Author: stankiewicz

runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java

@@ -46,7 +46,8 @@ public interface DoFnRunner<InputT extends @Nullable Object, OutputT extends @Nu
       BoundedWindow window,
       Instant timestamp,
       Instant outputTimestamp,
-      TimeDomain timeDomain);
+      TimeDomain timeDomain,
+      boolean causedByDrain);
 
   /**
    * Calls a {@link DoFn DoFn's} {@link DoFn.FinishBundle @FinishBundle} method and performs

runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java

@@ -89,8 +89,10 @@ public <KeyT> void onTimer(
       BoundedWindow window,
       Instant timestamp,
       Instant outputTimestamp,
-      TimeDomain timeDomain) {
-    doFnRunner.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain);
+      TimeDomain timeDomain,
+      boolean causedByDrain) {
+    doFnRunner.onTimer(
+        timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, causedByDrain);
   }
 
   @Override

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

@@ -396,6 +396,11 @@ public PaneInfo pane() {
       return element.getRecordOffset();
     }
 
+    @Override
+    public boolean causedByDrain() {
+      return false;
+    }
+
     @Override
     public PipelineOptions getPipelineOptions() {
       return pipelineOptions;

runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java

@@ -200,11 +200,13 @@ public <KeyT> void onTimer(
       BoundedWindow window,
       Instant timestamp,
       Instant outputTimestamp,
-      TimeDomain timeDomain) {
+      TimeDomain timeDomain,
+      boolean 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);
   }
 
@@ -399,6 +401,11 @@ public InputT element() {
       return elem.getValue();
     }
 
+    @Override
+    public boolean causedByDrain() {
+      return elem.causedByDrain();
+    }
+
     @Override
     public <T> T sideInput(PCollectionView<T> view) {
       checkNotNull(view, "View passed to sideInput cannot be null");
@@ -702,6 +709,7 @@ private class OnTimerArgumentProvider<KeyT> extends DoFn<InputT, OutputT>.OnTime
     private final TimeDomain timeDomain;
     private final String timerId;
     private final KeyT key;
+    private final Boolean causedByDrain;
     private final OutputBuilderSupplier builderSupplier;
 
     /** Lazily initialized; should only be accessed via {@link #getNamespace()}. */
@@ -727,28 +735,36 @@ private OnTimerArgumentProvider(
         BoundedWindow window,
         Instant fireTimestamp,
         Instant timestamp,
-        TimeDomain timeDomain) {
+        TimeDomain timeDomain,
+        boolean 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 boolean causedByDrain() {
+      return causedByDrain;
+    }
+
     @Override
     public Instant fireTimestamp() {
       return fireTimestamp;

runners/core-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java

@@ -116,7 +116,8 @@ 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, false);
   }
 
   @Override

runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java

@@ -48,7 +48,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
@@ -208,7 +209,8 @@ public <KeyT> void onTimer(
       BoundedWindow window,
       Instant timestamp,
       Instant outputTimestamp,
-      TimeDomain timeDomain) {
+      TimeDomain timeDomain,
+      boolean causedByDrain) {
 
     if (timerId.equals(SORT_FLUSH_TIMER)) {
       onSortFlushTimer(window, stepContext.timerInternals().currentInputWatermarkTime());
@@ -232,7 +234,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);
       }
     }
   }

runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java

@@ -140,7 +140,8 @@ public void testOnTimerExceptionsWrappedAsUserCodeException() {
         GlobalWindow.INSTANCE,
         new Instant(0),
         new Instant(0),
-        TimeDomain.EVENT_TIME);
+        TimeDomain.EVENT_TIME,
+        false);
   }
 
   /**
@@ -265,7 +266,8 @@ public void testOnTimerCalled() {
         GlobalWindow.INSTANCE,
         currentTime.plus(offset),
         currentTime.plus(offset),
-        TimeDomain.EVENT_TIME);
+        TimeDomain.EVENT_TIME,
+        false);
 
     assertThat(
         fn.onTimerInvocations,
@@ -276,7 +278,8 @@ public void testOnTimerCalled() {
                 StateNamespaces.window(windowFn.windowCoder(), GlobalWindow.INSTANCE),
                 currentTime.plus(offset),
                 currentTime.plus(offset),
-                TimeDomain.EVENT_TIME)));
+                TimeDomain.EVENT_TIME,
+                false)));
   }
 
   /**
@@ -592,7 +595,8 @@ public void testOnTimerAllowedSkew() {
         GlobalWindow.INSTANCE,
         new Instant(0),
         new Instant(0),
-        TimeDomain.EVENT_TIME);
+        TimeDomain.EVENT_TIME,
+        false);
   }
 
   @Test
@@ -624,7 +628,8 @@ public void testOnTimerNoSkew() {
                   GlobalWindow.INSTANCE,
                   new Instant(0),
                   new Instant(0),
-                  TimeDomain.EVENT_TIME);
+                  TimeDomain.EVENT_TIME,
+                  false);
             });
 
     assertThat(exception.getCause(), isA(IllegalArgumentException.class));
@@ -702,7 +707,7 @@ public void onTimer(OnTimerContext context) {
               context.fireTimestamp(),
               context.timestamp(),
               context.timeDomain(),
-              false));
+              context.causedByDrain()));
     }
   }
 

runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java

@@ -352,7 +352,8 @@ public <KeyT> void onTimer(
         BoundedWindow window,
         Instant timestamp,
         Instant outputTimestamp,
-        TimeDomain timeDomain) {
+        TimeDomain timeDomain,
+        boolean causedByDrain) {
       firedTimers.add(
           TimerData.of(
               timerId,
@@ -361,7 +362,7 @@ public <KeyT> void onTimer(
               timestamp,
               outputTimestamp,
               timeDomain,
-              false));
+              causedByDrain));
     }
 
     @Override

runners/core-java/src/test/java/org/apache/beam/runners/core/StatefulDoFnRunnerTest.java

@@ -468,7 +468,8 @@ private static void advanceInputWatermark(
           window,
           timer.getTimestamp(),
           timer.getOutputTimestamp(),
-          timer.getDomain());
+          timer.getDomain(),
+          timer.causedByDrain());
     }
   }
 

runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java

@@ -73,10 +73,19 @@ public <KeyT> void onTimer(
       final BoundedWindow window,
       final Instant timestamp,
       final Instant outputTimestamp,
-      final TimeDomain timeDomain) {
+      final TimeDomain timeDomain,
+      boolean 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);
     }