Examples with DataOutputViewStreamWrapper org.apache.flink.core.memory.DataOutputViewStreamWrapper used on opensource projects

Example 76 with DataOutputViewStreamWrapper

use of org.apache.flink.core.memory.DataOutputViewStreamWrapper in project flink by apache.

the class HeapInternalTimerServiceTest method testSnapshotAndRebalancingRestore.

/**
	 * This test checks whether timers are assigned to correct key groups
	 * and whether snapshot/restore respects key groups.
	 */
@Test
public void testSnapshotAndRebalancingRestore() throws Exception {
    @SuppressWarnings("unchecked") Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);
    TestKeyContext keyContext = new TestKeyContext();
    TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
    HeapInternalTimerService<Integer, String> timerService = createTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, maxParallelism);
    int midpoint = testKeyGroupRange.getStartKeyGroup() + (testKeyGroupRange.getEndKeyGroup() - testKeyGroupRange.getStartKeyGroup()) / 2;
    // get two sub key-ranges so that we can restore two ranges separately
    KeyGroupRange subKeyGroupRange1 = new KeyGroupRange(testKeyGroupRange.getStartKeyGroup(), midpoint);
    KeyGroupRange subKeyGroupRange2 = new KeyGroupRange(midpoint + 1, testKeyGroupRange.getEndKeyGroup());
    // get two different keys, one per sub range
    int key1 = getKeyInKeyGroupRange(subKeyGroupRange1, maxParallelism);
    int key2 = getKeyInKeyGroupRange(subKeyGroupRange2, maxParallelism);
    keyContext.setCurrentKey(key1);
    timerService.registerProcessingTimeTimer("ciao", 10);
    timerService.registerEventTimeTimer("hello", 10);
    keyContext.setCurrentKey(key2);
    timerService.registerEventTimeTimer("ciao", 10);
    timerService.registerProcessingTimeTimer("hello", 10);
    assertEquals(2, timerService.numProcessingTimeTimers());
    assertEquals(1, timerService.numProcessingTimeTimers("hello"));
    assertEquals(1, timerService.numProcessingTimeTimers("ciao"));
    assertEquals(2, timerService.numEventTimeTimers());
    assertEquals(1, timerService.numEventTimeTimers("hello"));
    assertEquals(1, timerService.numEventTimeTimers("ciao"));
    // one map per sub key-group range
    Map<Integer, byte[]> snapshot1 = new HashMap<>();
    Map<Integer, byte[]> snapshot2 = new HashMap<>();
    for (Integer keyGroupIndex : testKeyGroupRange) {
        ByteArrayOutputStream outStream = new ByteArrayOutputStream();
        timerService.snapshotTimersForKeyGroup(new DataOutputViewStreamWrapper(outStream), keyGroupIndex);
        outStream.close();
        if (subKeyGroupRange1.contains(keyGroupIndex)) {
            snapshot1.put(keyGroupIndex, outStream.toByteArray());
        } else if (subKeyGroupRange2.contains(keyGroupIndex)) {
            snapshot2.put(keyGroupIndex, outStream.toByteArray());
        } else {
            throw new IllegalStateException("Key-Group index doesn't belong to any sub range.");
        }
    }
    // from now on we need everything twice. once per sub key-group range
    @SuppressWarnings("unchecked") Triggerable<Integer, String> mockTriggerable1 = mock(Triggerable.class);
    @SuppressWarnings("unchecked") Triggerable<Integer, String> mockTriggerable2 = mock(Triggerable.class);
    TestKeyContext keyContext1 = new TestKeyContext();
    TestKeyContext keyContext2 = new TestKeyContext();
    TestProcessingTimeService processingTimeService1 = new TestProcessingTimeService();
    TestProcessingTimeService processingTimeService2 = new TestProcessingTimeService();
    HeapInternalTimerService<Integer, String> timerService1 = restoreTimerService(snapshot1, mockTriggerable1, keyContext1, processingTimeService1, subKeyGroupRange1, maxParallelism);
    HeapInternalTimerService<Integer, String> timerService2 = restoreTimerService(snapshot2, mockTriggerable2, keyContext2, processingTimeService2, subKeyGroupRange2, maxParallelism);
    processingTimeService1.setCurrentTime(10);
    timerService1.advanceWatermark(10);
    verify(mockTriggerable1, times(1)).onProcessingTime(anyInternalTimer());
    verify(mockTriggerable1, times(1)).onProcessingTime(eq(new InternalTimer<>(10, key1, "ciao")));
    verify(mockTriggerable1, never()).onProcessingTime(eq(new InternalTimer<>(10, key2, "hello")));
    verify(mockTriggerable1, times(1)).onEventTime(anyInternalTimer());
    verify(mockTriggerable1, times(1)).onEventTime(eq(new InternalTimer<>(10, key1, "hello")));
    verify(mockTriggerable1, never()).onEventTime(eq(new InternalTimer<>(10, key2, "ciao")));
    assertEquals(0, timerService1.numEventTimeTimers());
    processingTimeService2.setCurrentTime(10);
    timerService2.advanceWatermark(10);
    verify(mockTriggerable2, times(1)).onProcessingTime(anyInternalTimer());
    verify(mockTriggerable2, never()).onProcessingTime(eq(new InternalTimer<>(10, key1, "ciao")));
    verify(mockTriggerable2, times(1)).onProcessingTime(eq(new InternalTimer<>(10, key2, "hello")));
    verify(mockTriggerable2, times(1)).onEventTime(anyInternalTimer());
    verify(mockTriggerable2, never()).onEventTime(eq(new InternalTimer<>(10, key1, "hello")));
    verify(mockTriggerable2, times(1)).onEventTime(eq(new InternalTimer<>(10, key2, "ciao")));
    assertEquals(0, timerService2.numEventTimeTimers());
}

Example 77 with DataOutputViewStreamWrapper

use of org.apache.flink.core.memory.DataOutputViewStreamWrapper in project flink by apache.

the class HeapKeyedStateBackend method snapshot.

@Override
@SuppressWarnings("unchecked")
public RunnableFuture<KeyGroupsStateHandle> snapshot(final long checkpointId, final long timestamp, final CheckpointStreamFactory streamFactory, CheckpointOptions checkpointOptions) throws Exception {
    if (!hasRegisteredState()) {
        return DoneFuture.nullValue();
    }
    long syncStartTime = System.currentTimeMillis();
    Preconditions.checkState(stateTables.size() <= Short.MAX_VALUE, "Too many KV-States: " + stateTables.size() + ". Currently at most " + Short.MAX_VALUE + " states are supported");
    List<KeyedBackendSerializationProxy.StateMetaInfo<?, ?>> metaInfoProxyList = new ArrayList<>(stateTables.size());
    final Map<String, Integer> kVStateToId = new HashMap<>(stateTables.size());
    final Map<StateTable<K, ?, ?>, StateTableSnapshot> cowStateStableSnapshots = new HashedMap(stateTables.size());
    for (Map.Entry<String, StateTable<K, ?, ?>> kvState : stateTables.entrySet()) {
        RegisteredBackendStateMetaInfo<?, ?> metaInfo = kvState.getValue().getMetaInfo();
        KeyedBackendSerializationProxy.StateMetaInfo<?, ?> metaInfoProxy = new KeyedBackendSerializationProxy.StateMetaInfo(metaInfo.getStateType(), metaInfo.getName(), metaInfo.getNamespaceSerializer(), metaInfo.getStateSerializer());
        metaInfoProxyList.add(metaInfoProxy);
        kVStateToId.put(kvState.getKey(), kVStateToId.size());
        StateTable<K, ?, ?> stateTable = kvState.getValue();
        if (null != stateTable) {
            cowStateStableSnapshots.put(stateTable, stateTable.createSnapshot());
        }
    }
    final KeyedBackendSerializationProxy serializationProxy = new KeyedBackendSerializationProxy(keySerializer, metaInfoProxyList);
    //--------------------------------------------------- this becomes the end of sync part
    // implementation of the async IO operation, based on FutureTask
    final AbstractAsyncIOCallable<KeyGroupsStateHandle, CheckpointStreamFactory.CheckpointStateOutputStream> ioCallable = new AbstractAsyncIOCallable<KeyGroupsStateHandle, CheckpointStreamFactory.CheckpointStateOutputStream>() {

        AtomicBoolean open = new AtomicBoolean(false);

        @Override
        public CheckpointStreamFactory.CheckpointStateOutputStream openIOHandle() throws Exception {
            if (open.compareAndSet(false, true)) {
                CheckpointStreamFactory.CheckpointStateOutputStream stream = streamFactory.createCheckpointStateOutputStream(checkpointId, timestamp);
                try {
                    cancelStreamRegistry.registerClosable(stream);
                    return stream;
                } catch (Exception ex) {
                    open.set(false);
                    throw ex;
                }
            } else {
                throw new IOException("Operation already opened.");
            }
        }

        @Override
        public KeyGroupsStateHandle performOperation() throws Exception {
            long asyncStartTime = System.currentTimeMillis();
            CheckpointStreamFactory.CheckpointStateOutputStream stream = getIoHandle();
            DataOutputViewStreamWrapper outView = new DataOutputViewStreamWrapper(stream);
            serializationProxy.write(outView);
            long[] keyGroupRangeOffsets = new long[keyGroupRange.getNumberOfKeyGroups()];
            for (int keyGroupPos = 0; keyGroupPos < keyGroupRange.getNumberOfKeyGroups(); ++keyGroupPos) {
                int keyGroupId = keyGroupRange.getKeyGroupId(keyGroupPos);
                keyGroupRangeOffsets[keyGroupPos] = stream.getPos();
                outView.writeInt(keyGroupId);
                for (Map.Entry<String, StateTable<K, ?, ?>> kvState : stateTables.entrySet()) {
                    outView.writeShort(kVStateToId.get(kvState.getKey()));
                    cowStateStableSnapshots.get(kvState.getValue()).writeMappingsInKeyGroup(outView, keyGroupId);
                }
            }
            if (open.compareAndSet(true, false)) {
                StreamStateHandle streamStateHandle = stream.closeAndGetHandle();
                KeyGroupRangeOffsets offsets = new KeyGroupRangeOffsets(keyGroupRange, keyGroupRangeOffsets);
                final KeyGroupsStateHandle keyGroupsStateHandle = new KeyGroupsStateHandle(offsets, streamStateHandle);
                if (asynchronousSnapshots) {
                    LOG.info("Heap backend snapshot ({}, asynchronous part) in thread {} took {} ms.", streamFactory, Thread.currentThread(), (System.currentTimeMillis() - asyncStartTime));
                }
                return keyGroupsStateHandle;
            } else {
                throw new IOException("Checkpoint stream already closed.");
            }
        }

        @Override
        public void done(boolean canceled) {
            if (open.compareAndSet(true, false)) {
                CheckpointStreamFactory.CheckpointStateOutputStream stream = getIoHandle();
                if (null != stream) {
                    cancelStreamRegistry.unregisterClosable(stream);
                    IOUtils.closeQuietly(stream);
                }
            }
            for (StateTableSnapshot snapshot : cowStateStableSnapshots.values()) {
                snapshot.release();
            }
        }
    };
    AsyncStoppableTaskWithCallback<KeyGroupsStateHandle> task = AsyncStoppableTaskWithCallback.from(ioCallable);
    if (!asynchronousSnapshots) {
        task.run();
    }
    LOG.info("Heap backend snapshot (" + streamFactory + ", synchronous part) in thread " + Thread.currentThread() + " took " + (System.currentTimeMillis() - syncStartTime) + " ms.");
    return task;
}

Example 78 with DataOutputViewStreamWrapper

use of org.apache.flink.core.memory.DataOutputViewStreamWrapper in project flink by apache.

the class AbstractPythonStreamAggregateOperator method open.

@Override
@SuppressWarnings("unchecked")
public void open() throws Exception {
    bais = new ByteArrayInputStreamWithPos();
    baisWrapper = new DataInputViewStreamWrapper(bais);
    baos = new ByteArrayOutputStreamWithPos();
    baosWrapper = new DataOutputViewStreamWrapper(baos);
    userDefinedFunctionInputType = createUserDefinedFunctionInputType();
    udfInputTypeSerializer = PythonTypeUtils.toInternalSerializer(userDefinedFunctionInputType);
    userDefinedFunctionOutputType = createUserDefinedFunctionOutputType();
    udfOutputTypeSerializer = PythonTypeUtils.toInternalSerializer(userDefinedFunctionOutputType);
    rowDataWrapper = new StreamRecordRowDataWrappingCollector(output);
    super.open();
    configJobOptions();
}

Example 79 with DataOutputViewStreamWrapper

use of org.apache.flink.core.memory.DataOutputViewStreamWrapper in project flink by apache.

the class AbstractStatelessFunctionOperator method open.

@Override
public void open() throws Exception {
    forwardedInputQueue = new LinkedList<>();
    bais = new ByteArrayInputStreamWithPos();
    baisWrapper = new DataInputViewStreamWrapper(bais);
    baos = new ByteArrayOutputStreamWithPos();
    baosWrapper = new DataOutputViewStreamWrapper(baos);
    super.open();
}

Example 80 with DataOutputViewStreamWrapper

use of org.apache.flink.core.memory.DataOutputViewStreamWrapper in project flink by apache.

the class BatchArrowPythonOverWindowAggregateFunctionOperator method open.

@Override
public void open() throws Exception {
    super.open();
    forwardedInputSerializer = new RowDataSerializer(inputType);
    this.lastKeyDataStartPos = 0;
    windowBoundaryWithDataBaos = new ByteArrayOutputStreamWithPos();
    windowBoundaryWithDataWrapper = new DataOutputViewStreamWrapper(windowBoundaryWithDataBaos);
    boundedRangeWindowBoundaries = new ArrayList<>(lowerBoundary.length);
    boundedRangeWindowIndex = new ArrayList<>();
    for (int i = 0; i < lowerBoundary.length; i++) {
        // range window with bounded preceding or bounded following
        if (isRangeWindows[i] && (lowerBoundary[i] != Long.MIN_VALUE || upperBoundary[i] != Long.MAX_VALUE)) {
            boundedRangeWindowIndex.add(i);
            boundedRangeWindowBoundaries.add(new ArrayList<>());
        }
    }
}