org.apache.beam.sdk.state.TimeDomain Java Examples

@Override
public <W extends BoundedWindow> void setStateCleanupTimer(
    String timerId,
    W window,
    Coder<W> windowCoder,
    Instant cleanupTime,
    Instant cleanupOutputTimestamp) {
  timerInternals()
      .setTimer(
          StateNamespaces.window(windowCoder, window),
          timerId,
          "",
          cleanupTime,
          cleanupOutputTimestamp,
          TimeDomain.EVENT_TIME);
}
 

@Override
public org.apache.beam.sdk.state.Timer timer(String timerId) {
  checkState(
      currentElement.getValue() instanceof KV,
      "Accessing timer in unkeyed context. Current element is not a KV: %s.",
      currentElement.getValue());

  // For the initial timestamps we pass in the current elements timestamp for the hold timestamp
  // and the current element's timestamp which will be used for the fire timestamp if this
  // timer is in the EVENT time domain.
  TimeDomain timeDomain =
      translateTimeDomain(parDoPayload.getTimerFamilySpecsMap().get(timerId).getTimeDomain());
  return new FnApiTimer(
      timerId,
      ((KV) currentElement.getValue()).getKey(),
      "",
      currentWindow,
      currentElement.getTimestamp(),
      currentElement.getTimestamp(),
      currentElement.getPane(),
      timeDomain);
}
 

@Override
public BufferedElement decode(InputStream inStream) throws IOException {
  int firstByte = inStream.read();
  switch (firstByte) {
    case ELEMENT_MAGIC_BYTE:
      return new Element(elementCoder.decode(inStream));
    case TIMER_MAGIC_BYTE:
      return new Timer<>(
          STRING_CODER.decode(inStream),
          STRING_CODER.decode(inStream),
          key,
          windowCoder.decode(inStream),
          INSTANT_CODER.decode(inStream),
          INSTANT_CODER.decode(inStream),
          TimeDomain.values()[inStream.read()]);
    default:
      throw new IllegalStateException(
          "Unexpected byte while reading BufferedElement: " + firstByte);
  }
}
 

@Override
public void onElement(OnElementContext c) throws Exception {
  if (!endOfWindowReached(c)) {
    c.setTimer(c.window().maxTimestamp(), TimeDomain.EVENT_TIME);
  }

  if (!c.trigger().isMerging()) {
    // If merges can never happen, we just run the unfinished subtrigger
    c.trigger().firstUnfinishedSubTrigger().invokeOnElement(c);
  } else {
    // If merges can happen, we run for all subtriggers because they might be
    // de-activated or re-activated
    for (ExecutableTriggerStateMachine subTrigger : c.trigger().subTriggers()) {
      subTrigger.invokeOnElement(c);
    }
  }
}
 

@Override
public KeyedTimerData<K> decode(InputStream inStream) throws CoderException, IOException {
  // decode the timestamp first
  final Instant timestamp = INSTANT_CODER.decode(inStream);
  final String timerId = STRING_CODER.decode(inStream);
  final StateNamespace namespace =
      StateNamespaces.fromString(STRING_CODER.decode(inStream), windowCoder);
  final TimeDomain domain = TimeDomain.valueOf(STRING_CODER.decode(inStream));
  final TimerData timer = TimerData.of(timerId, namespace, timestamp, timestamp, domain);

  byte[] keyBytes = null;
  K key = null;
  if (keyCoder != null) {
    key = keyCoder.decode(inStream);

    final ByteArrayOutputStream baos = new ByteArrayOutputStream();
    try {
      keyCoder.encode(key, baos);
    } catch (IOException e) {
      throw new RuntimeException("Could not encode key: " + key, e);
    }
    keyBytes = baos.toByteArray();
  }

  return new KeyedTimerData(keyBytes, key, timer);
}
 

private void startBundle() {
  // Register as a consumer for each timer.
  timerHandlers = new HashMap<>();
  for (Map.Entry<String, KV<TimeDomain, Coder<Timer<Object>>>> timerFamilyInfo :
      timerFamilyInfos.entrySet()) {
    String localName = timerFamilyInfo.getKey();
    TimeDomain timeDomain = timerFamilyInfo.getValue().getKey();
    Coder<Timer<Object>> timerCoder = timerFamilyInfo.getValue().getValue();
    timerHandlers.put(
        localName,
        beamFnTimerClient.register(
            LogicalEndpoint.timer(processBundleInstructionId.get(), pTransformId, localName),
            timerCoder,
            (FnDataReceiver<Timer<Object>>) timer -> processTimer(localName, timeDomain, timer)));
  }

  doFnInvoker.invokeStartBundle(startBundleArgumentProvider);
}
 

@Test
public void testTimerDataCoder() throws Exception {
  CoderProperties.coderDecodeEncodeEqual(
      TimerDataCoderV2.of(GlobalWindow.Coder.INSTANCE),
      TimerData.of(
          "arbitrary-id",
          StateNamespaces.global(),
          new Instant(0),
          new Instant(0),
          TimeDomain.EVENT_TIME));

  Coder<IntervalWindow> windowCoder = IntervalWindow.getCoder();
  CoderProperties.coderDecodeEncodeEqual(
      TimerDataCoderV2.of(windowCoder),
      TimerData.of(
          "another-id",
          StateNamespaces.window(
              windowCoder, new IntervalWindow(new Instant(0), new Instant(100))),
          new Instant(99),
          new Instant(99),
          TimeDomain.PROCESSING_TIME));
}
 

/**
 * Returns when the next timer in the given time domain will fire, or {@code null} if there are no
 * timers scheduled in that time domain.
 */
@Nullable
public Instant getNextTimer(TimeDomain domain) {
  try {
    switch (domain) {
      case EVENT_TIME:
        return watermarkTimers.first().getTimestamp();
      case PROCESSING_TIME:
        return processingTimers.first().getTimestamp();
      case SYNCHRONIZED_PROCESSING_TIME:
        return synchronizedProcessingTimers.first().getTimestamp();
      default:
        throw new IllegalArgumentException("Unexpected time domain: " + domain);
    }
  } catch (NoSuchElementException exc) {
    return null;
  }
}
 

/**
 * Setup timer for flush time @{code flush}. The time is adjusted to respect allowed lateness and
 * window garbage collection time. Setup watermark hold for the flush time.
 *
 * <p>Note that this is equivalent to {@link org.apache.beam.sdk.state.Timer#withOutputTimestamp}
 * and should be reworked to use that feature once that is stable.
 */
private void setupFlushTimerAndWatermarkHold(
    StateNamespace namespace, BoundedWindow window, Instant flush) {
  Instant flushWithLateness = flush.plus(windowingStrategy.getAllowedLateness());
  Instant windowGcTime = window.maxTimestamp().plus(windowingStrategy.getAllowedLateness());
  if (flushWithLateness.isAfter(windowGcTime)) {
    flushWithLateness = windowGcTime;
  }
  WatermarkHoldState watermark = stepContext.stateInternals().state(namespace, watermarkHold);
  stepContext
      .timerInternals()
      .setTimer(
          namespace,
          SORT_FLUSH_TIMER,
          SORT_FLUSH_TIMER,
          flushWithLateness,
          flush,
          TimeDomain.EVENT_TIME);
  watermark.clear();
  watermark.add(flush);
}
 

private synchronized Map<StructuralKey<?>, List<TimerData>> extractFiredDomainTimers(
    TimeDomain domain, Instant firingTime) {
  Map<StructuralKey<?>, List<TimerData>> firedTimers;
  switch (domain) {
    case PROCESSING_TIME:
      firedTimers = extractFiredTimers(firingTime, processingTimers);
      break;
    case SYNCHRONIZED_PROCESSING_TIME:
      firedTimers =
          extractFiredTimers(
              INSTANT_ORDERING.min(firingTime, earliestHold.get()),
              synchronizedProcessingTimers);
      break;
    default:
      throw new IllegalArgumentException(
          "Called getFiredTimers on a Synchronized Processing Time watermark"
              + " and gave a non-processing time domain "
              + domain);
  }
  for (Map.Entry<StructuralKey<?>, ? extends Collection<TimerData>> firedTimer :
      firedTimers.entrySet()) {
    pendingTimers.addAll(firedTimer.getValue());
  }
  return firedTimers;
}
 

private void cancelEndOfWindowAndGarbageCollectionTimers(
    ReduceFn<?, ?, ?, W>.Context directContext) {
  WindowTracing.debug(
      "ReduceFnRunner.cancelEndOfWindowAndGarbageCollectionTimers: Deleting timers for "
          + "key:{}; window:{} where inputWatermark:{}; outputWatermark:{}",
      key,
      directContext.window(),
      timerInternals.currentInputWatermarkTime(),
      timerInternals.currentOutputWatermarkTime());
  Instant eow = directContext.window().maxTimestamp();
  directContext.timers().deleteTimer(eow, TimeDomain.EVENT_TIME);
  Instant gc = LateDataUtils.garbageCollectionTime(directContext.window(), windowingStrategy);
  if (gc.isAfter(eow)) {
    directContext.timers().deleteTimer(gc, TimeDomain.EVENT_TIME);
  }
}
 

/**
 * Schedule a timer to garbage collect the window.
 *
 * <p>The timer:
 *
 * <ul>
 *   <li>...must be fired strictly after the expiration of the window.
 *   <li>...should be as close to the expiration as possible, to have a timely output of remaining
 *       buffered data, and GC.
 * </ul>
 */
private void scheduleGarbageCollectionTimer(ReduceFn<?, ?, ?, W>.Context directContext) {
  Instant inputWM = timerInternals.currentInputWatermarkTime();
  Instant gcTime = LateDataUtils.garbageCollectionTime(directContext.window(), windowingStrategy);
  WindowTracing.trace(
      "ReduceFnRunner.scheduleGarbageCollectionTimer: Scheduling at {} for "
          + "key:{}; window:{} where inputWatermark:{}; outputWatermark:{}",
      gcTime,
      key,
      directContext.window(),
      inputWM,
      timerInternals.currentOutputWatermarkTime());
  checkState(
      !gcTime.isAfter(BoundedWindow.TIMESTAMP_MAX_VALUE),
      "Timer %s is beyond end-of-time",
      gcTime);
  directContext.timers().setTimer(gcTime, TimeDomain.EVENT_TIME);
}
 

@Test
public void testCompareByNamespace() {
  Instant timestamp = new Instant(100);
  IntervalWindow firstWindow = new IntervalWindow(new Instant(0), timestamp);
  IntervalWindow secondWindow = new IntervalWindow(timestamp, new Instant(200));
  Coder<IntervalWindow> windowCoder = IntervalWindow.getCoder();

  StateNamespace firstWindowNs = StateNamespaces.window(windowCoder, firstWindow);
  StateNamespace secondWindowNs = StateNamespaces.window(windowCoder, secondWindow);

  TimerData secondEventTime =
      TimerData.of(firstWindowNs, timestamp, timestamp, TimeDomain.EVENT_TIME);
  TimerData thirdEventTime =
      TimerData.of(secondWindowNs, timestamp, timestamp, TimeDomain.EVENT_TIME);

  assertThat(secondEventTime, lessThan(thirdEventTime));
}
 

@Test
public void testWindowParamOnTimer() throws Exception {
  final String timerId = "some-timer-id";
  final String timerDeclarationId = TimerDeclaration.PREFIX + timerId;

  DoFnSignature sig =
      DoFnSignatures.getSignature(
          new DoFn<String, String>() {
            @TimerId(timerId)
            private final TimerSpec myfield1 = TimerSpecs.timer(TimeDomain.EVENT_TIME);

            @ProcessElement
            public void process(ProcessContext c) {}

            @OnTimer(timerId)
            public void onTimer(BoundedWindow w) {}
          }.getClass());

  assertThat(sig.onTimerMethods().get(timerDeclarationId).extraParameters().size(), equalTo(1));
  assertThat(
      sig.onTimerMethods().get(timerDeclarationId).extraParameters().get(0),
      instanceOf(WindowParameter.class));
}
 

static <K> KeyedTimerData<K> toKeyedTimerData(
    TimerKey<K> timerKey, long timestamp, TimeDomain domain, Coder<K> keyCoder) {
  byte[] keyBytes = null;
  if (keyCoder != null && timerKey.key != null) {
    final ByteArrayOutputStream baos = new ByteArrayOutputStream();
    try {
      keyCoder.encode(timerKey.key, baos);
    } catch (IOException e) {
      throw new RuntimeException("Could not encode key: " + timerKey.key, e);
    }
    keyBytes = baos.toByteArray();
  }

  return new KeyedTimerData<K>(
      keyBytes,
      timerKey.key,
      TimerInternals.TimerData.of(
          timerKey.timerId,
          timerKey.stateNamespace,
          new Instant(timestamp),
          new Instant(timestamp),
          domain));
}
 

@Test
public void testSimpleTimerWithContext() throws Exception {
  DoFnSignature sig =
      DoFnSignatures.getSignature(
          new DoFn<KV<String, Integer>, Long>() {
            @TimerId("foo")
            private final TimerSpec bizzle = TimerSpecs.timer(TimeDomain.EVENT_TIME);

            @ProcessElement
            public void foo(ProcessContext context) {}

            @OnTimer("foo")
            public void onFoo(OnTimerContext c) {}
          }.getClass());

  final String timerDeclarationId = TimerDeclaration.PREFIX + "foo";
  assertThat(sig.timerDeclarations().size(), equalTo(1));
  DoFnSignature.TimerDeclaration decl = sig.timerDeclarations().get(timerDeclarationId);

  assertThat(decl.id(), equalTo(timerDeclarationId));
  assertThat(decl.field().getName(), equalTo("bizzle"));

  assertThat(
      sig.onTimerMethods().get(timerDeclarationId).extraParameters().get(0),
      equalTo((Parameter) Parameter.onTimerContext()));
}
 

@Test
public void testSimpleTimerIdAnonymousDoFn() throws Exception {
  DoFnSignature sig =
      DoFnSignatures.getSignature(
          new DoFn<KV<String, Integer>, Long>() {
            @TimerId("foo")
            private final TimerSpec bizzle = TimerSpecs.timer(TimeDomain.EVENT_TIME);

            @ProcessElement
            public void foo(ProcessContext context) {}

            @OnTimer("foo")
            public void onFoo() {}
          }.getClass());

  final String timerDeclarationId = TimerDeclaration.PREFIX + "foo";
  assertThat(sig.timerDeclarations().size(), equalTo(1));
  DoFnSignature.TimerDeclaration decl = sig.timerDeclarations().get(timerDeclarationId);

  assertThat(decl.id(), equalTo(timerDeclarationId));
  assertThat(decl.field().getName(), equalTo("bizzle"));
}
 

private TimerData timerData(IntervalWindow window, Instant timestamp, Timer.Type type) {
  return TimerData.of(
      StateNamespaces.window(IntervalWindow.getCoder(), window),
      timestamp,
      timestamp,
      type == Windmill.Timer.Type.WATERMARK ? TimeDomain.EVENT_TIME : TimeDomain.PROCESSING_TIME);
}
 

@Override
public TimerData decode(InputStream inStream) throws CoderException, IOException {
  String timerId = STRING_CODER.decode(inStream);
  StateNamespace namespace =
      StateNamespaces.fromString(STRING_CODER.decode(inStream), windowCoder);
  Instant timestamp = INSTANT_CODER.decode(inStream);
  TimeDomain domain = TimeDomain.valueOf(STRING_CODER.decode(inStream));
  return TimerData.of(timerId, namespace, timestamp, timestamp, domain);
}
 

private NavigableSet<TimerData> timersForDomain(TimeDomain domain) {
  switch (domain) {
    case EVENT_TIME:
      return watermarkTimers;
    case PROCESSING_TIME:
      return processingTimers;
    case SYNCHRONIZED_PROCESSING_TIME:
      return synchronizedProcessingTimers;
    default:
      throw new IllegalArgumentException("Unexpected time domain: " + domain);
  }
}
 

@Test
public void testCompareEqual() {
  Instant timestamp = new Instant(100);
  StateNamespace namespace = StateNamespaces.global();
  TimerData timer = TimerData.of("id", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME);

  assertThat(
      timer,
      comparesEqualTo(
          TimerData.of("id", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME)));
}
 

@Override
public <KeyT> void onTimer(
    final String timerId,
    final String timerFamilyId,
    KeyT key,
    final BoundedWindow window,
    final Instant timestamp,
    final Instant outputTimestamp,
    final TimeDomain timeDomain) {
  try (Closeable ignored = MetricsEnvironment.scopedMetricsContainer(metricsContainer())) {
    delegate.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain);
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
}
 

/** Verifies that the time domain of this timer is acceptable for absolute timers. */
private void verifyAbsoluteTimeDomain() {
  if (!TimeDomain.EVENT_TIME.equals(spec.getTimeDomain())) {
    throw new IllegalStateException(
        "Cannot only set relative timers in processing time domain." + " Use #setRelative()");
  }
}
 

@Test
public void testOnTimerExceptionsWrappedAsUserCodeException() {
  ThrowingDoFn fn = new ThrowingDoFn();
  DoFnRunner<String, String> runner =
      new SimpleDoFnRunner<>(
          null,
          fn,
          NullSideInputReader.empty(),
          null,
          null,
          Collections.emptyList(),
          mockStepContext,
          null,
          Collections.emptyMap(),
          WindowingStrategy.of(new GlobalWindows()),
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  thrown.expect(UserCodeException.class);
  thrown.expectCause(is(fn.exceptionToThrow));

  runner.onTimer(
      TimerDeclaration.PREFIX + ThrowingDoFn.TIMER_ID,
      "",
      null,
      GlobalWindow.INSTANCE,
      new Instant(0),
      new Instant(0),
      TimeDomain.EVENT_TIME);
}
 

@Test
public void timerUpdateBuilderBuildAddsAllAddedTimers() {
  TimerData set =
      TimerData.of(
          StateNamespaces.global(), new Instant(10L), new Instant(10L), TimeDomain.EVENT_TIME);
  TimerData deleted =
      TimerData.of(
          StateNamespaces.global(),
          new Instant(24L),
          new Instant(24L),
          TimeDomain.PROCESSING_TIME);
  TimerData completedOne =
      TimerData.of(
          StateNamespaces.global(),
          new Instant(1024L),
          new Instant(1024L),
          TimeDomain.SYNCHRONIZED_PROCESSING_TIME);
  TimerData completedTwo =
      TimerData.of(
          StateNamespaces.global(),
          new Instant(2048L),
          new Instant(2048L),
          TimeDomain.EVENT_TIME);

  TimerUpdate update =
      TimerUpdate.builder(StructuralKey.of("foo", StringUtf8Coder.of()))
          .withCompletedTimers(ImmutableList.of(completedOne, completedTwo))
          .setTimer(set)
          .deletedTimer(deleted)
          .build();

  assertThat(update.getCompletedTimers(), containsInAnyOrder(completedOne, completedTwo));
  assertThat(update.getSetTimers(), contains(set));
  assertThat(update.getDeletedTimers(), contains(deleted));
}
 

@Test
public void testDeletionIdempotent() throws Exception {
  InMemoryTimerInternals underTest = new InMemoryTimerInternals();
  Instant timestamp = new Instant(42);
  underTest.setTimer(NS1, ID1, ID1, timestamp, timestamp, TimeDomain.EVENT_TIME);
  underTest.deleteTimer(NS1, ID1, ID1);
  underTest.deleteTimer(NS1, ID1, ID1);
}
 

@Override
public void onElement(OnElementContext c) throws Exception {
  // We're interested in knowing when the input watermark passes the end of the window.
  // (It is possible this has already happened, in which case the timer will be fired
  // almost immediately).
  if (!endOfWindowReached(c)) {
    c.setTimer(c.window().maxTimestamp(), TimeDomain.EVENT_TIME);
  }
}
 

/** Construct a {@link TimerData} for the given parameters. */
public static TimerData of(
    String timerId,
    StateNamespace namespace,
    Instant timestamp,
    Instant outputTimestamp,
    TimeDomain domain) {
  return new AutoValue_TimerInternals_TimerData(
      timerId, "", namespace, timestamp, outputTimestamp, domain);
}
 

/**
 * Tests that when a processing time timer comes in after a window is expired but in the same
 * bundle it does not cause a spurious output.
 */
@Test
public void testCombiningAccumulatingProcessingTime() throws Exception {
  WindowingStrategy<?, IntervalWindow> strategy =
      WindowingStrategy.of((WindowFn<?, IntervalWindow>) FixedWindows.of(Duration.millis(100)))
          .withTimestampCombiner(TimestampCombiner.EARLIEST)
          .withMode(AccumulationMode.ACCUMULATING_FIRED_PANES)
          .withAllowedLateness(Duration.ZERO)
          .withTrigger(
              Repeatedly.forever(
                  AfterProcessingTime.pastFirstElementInPane().plusDelayOf(Duration.millis(10))));

  ReduceFnTester<Integer, Integer, IntervalWindow> tester =
      ReduceFnTester.combining(strategy, Sum.ofIntegers(), VarIntCoder.of());

  tester.advanceProcessingTime(new Instant(5000));
  injectElement(tester, 2); // processing timer @ 5000 + 10; EOW timer @ 100
  injectElement(tester, 5);

  tester.advanceInputWatermarkNoTimers(new Instant(100));
  tester.advanceProcessingTimeNoTimers(new Instant(5010));

  // Fires the GC/EOW timer at the same time as the processing time timer.
  tester.fireTimers(
      new IntervalWindow(new Instant(0), new Instant(100)),
      TimestampedValue.of(TimeDomain.EVENT_TIME, new Instant(100)),
      TimestampedValue.of(TimeDomain.PROCESSING_TIME, new Instant(5010)));

  assertThat(
      tester.extractOutput(),
      contains(
          isSingleWindowedValue(
              equalTo(7), 2, 0, 100, PaneInfo.createPane(true, true, Timing.ON_TIME, 0, 0))));
}
 

@Override
public <KeyT> void onTimer(
    String timerId,
    String timerFamilyId,
    KeyT key,
    BoundedWindow window,
    Instant timestamp,
    Instant outputTimestamp,
    TimeDomain timeDomain) {}