Examples with StateNamespace org.apache.beam.runners.core.StateNamespace used on opensource projects

Example 11 with StateNamespace

use of org.apache.beam.runners.core.StateNamespace in project beam by apache.

the class SamzaTimerInternalsFactoryTest method testRestore.

@Test
public void testRestore() throws Exception {
    final SamzaPipelineOptions pipelineOptions = PipelineOptionsFactory.create().as(SamzaPipelineOptions.class);
    KeyValueStore<ByteArray, StateValue<?>> store = createStore();
    final SamzaTimerInternalsFactory<String> timerInternalsFactory = createTimerInternalsFactory(null, "timer", pipelineOptions, store);
    final String key = "testKey";
    final StateNamespace nameSpace = StateNamespaces.global();
    final TimerInternals timerInternals = timerInternalsFactory.timerInternalsForKey(key);
    final TimerInternals.TimerData timer1 = TimerInternals.TimerData.of("timer1", nameSpace, new Instant(10), new Instant(10), TimeDomain.EVENT_TIME);
    timerInternals.setTimer(timer1);
    final TimerInternals.TimerData timer2 = TimerInternals.TimerData.of("timer2", nameSpace, new Instant(100), new Instant(100), TimeDomain.EVENT_TIME);
    timerInternals.setTimer(timer2);
    store.close();
    // restore by creating a new instance
    store = createStore();
    final SamzaTimerInternalsFactory<String> restoredFactory = createTimerInternalsFactory(null, "timer", pipelineOptions, store);
    restoredFactory.setInputWatermark(new Instant(150));
    Collection<KeyedTimerData<String>> readyTimers = restoredFactory.removeReadyTimers();
    assertEquals(2, readyTimers.size());
    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    StringUtf8Coder.of().encode(key, baos);
    byte[] keyBytes = baos.toByteArray();
    assertEquals(readyTimers, Arrays.asList(new KeyedTimerData<>(keyBytes, key, timer1), new KeyedTimerData<>(keyBytes, key, timer2)));
    store.close();
}

Example 12 with StateNamespace

use of org.apache.beam.runners.core.StateNamespace in project beam by apache.

the class SamzaTimerInternalsFactoryTest method testBufferSizeNotExceedingPipelineOptionValue.

@Test
public void testBufferSizeNotExceedingPipelineOptionValue() {
    final SamzaPipelineOptions pipelineOptions = PipelineOptionsFactory.create().as(SamzaPipelineOptions.class);
    pipelineOptions.setEventTimerBufferSize(2);
    final KeyValueStore<ByteArray, StateValue<?>> store = createStore();
    final SamzaTimerInternalsFactory<String> timerInternalsFactory = createTimerInternalsFactory(null, "timer", pipelineOptions, store);
    final StateNamespace nameSpace = StateNamespaces.global();
    final TimerInternals timerInternals = timerInternalsFactory.timerInternalsForKey("testKey");
    // timers in store are then timestamped from 0 - 4.
    for (int i = 0; i < 5; i++) {
        timerInternals.setTimer(nameSpace, "timer" + i, "", new Instant(i), new Instant(i), TimeDomain.EVENT_TIME);
    }
    // only two timers are supposed to be in event time buffer
    assertEquals(2, timerInternalsFactory.getEventTimeBuffer().size());
    store.close();
}

Example 13 with StateNamespace

use of org.apache.beam.runners.core.StateNamespace in project beam by apache.

the class SamzaTimerInternalsFactoryTest method testMaxExpiredEventTimersProcessAtOnce.

private void testMaxExpiredEventTimersProcessAtOnce(int totalNumberOfTimersInStore, int totalNumberOfExpiredTimers, int maxExpiredTimersToProcessOnce, int expectedExpiredTimersToProcess) {
    final SamzaPipelineOptions pipelineOptions = PipelineOptionsFactory.create().as(SamzaPipelineOptions.class);
    pipelineOptions.setMaxReadyTimersToProcessOnce(maxExpiredTimersToProcessOnce);
    final KeyValueStore<ByteArray, StateValue<?>> store = createStore();
    final SamzaTimerInternalsFactory<String> timerInternalsFactory = createTimerInternalsFactory(null, "timer", pipelineOptions, store);
    final StateNamespace nameSpace = StateNamespaces.global();
    final TimerInternals timerInternals = timerInternalsFactory.timerInternalsForKey("testKey");
    TimerInternals.TimerData timer;
    for (int i = 0; i < totalNumberOfTimersInStore; i++) {
        timer = TimerInternals.TimerData.of("timer" + i, nameSpace, new Instant(i), new Instant(i), TimeDomain.EVENT_TIME);
        timerInternals.setTimer(timer);
    }
    // Set the timestamp of the input watermark to be the value of totalNumberOfExpiredTimers
    // so that totalNumberOfExpiredTimers timers are expected be expired with respect to this
    // watermark.
    final Instant inputWatermark = new Instant(totalNumberOfExpiredTimers);
    timerInternalsFactory.setInputWatermark(inputWatermark);
    final Collection<KeyedTimerData<String>> readyTimers = timerInternalsFactory.removeReadyTimers();
    assertEquals(expectedExpiredTimersToProcess, readyTimers.size());
    store.close();
}

Example 14 with StateNamespace

use of org.apache.beam.runners.core.StateNamespace in project beam by apache.

the class DoFnOperator method fireTimer.

// allow overriding this in WindowDoFnOperator
protected void fireTimer(TimerData timerData) {
    LOG.debug("Firing timer: {} at {} with output time {}", timerData.getTimerId(), timerData.getTimestamp().getMillis(), timerData.getOutputTimestamp().getMillis());
    StateNamespace namespace = timerData.getNamespace();
    // This is a user timer, so namespace must be WindowNamespace
    checkArgument(namespace instanceof WindowNamespace);
    BoundedWindow window = ((WindowNamespace) namespace).getWindow();
    timerInternals.onFiredOrDeletedTimer(timerData);
    pushbackDoFnRunner.onTimer(timerData.getTimerId(), timerData.getTimerFamilyId(), keyedStateInternals.getKey(), window, timerData.getTimestamp(), timerData.getOutputTimestamp(), timerData.getDomain());
}

Example 15 with StateNamespace

use of org.apache.beam.runners.core.StateNamespace in project beam by apache.

the class DoFnOp method fireTimer.

private void fireTimer(KeyedTimerData<?> keyedTimerData) {
    final TimerInternals.TimerData timer = keyedTimerData.getTimerData();
    LOG.debug("Firing timer {}", timer);
    final StateNamespace namespace = timer.getNamespace();
    // NOTE: not sure why this is safe, but DoFnOperator makes this assumption
    final BoundedWindow window = ((StateNamespaces.WindowNamespace) namespace).getWindow();
    fnRunner.onTimer(timer.getTimerId(), timer.getTimerFamilyId(), keyedTimerData.getKey(), window, timer.getTimestamp(), timer.getOutputTimestamp(), timer.getDomain());
}