Examples with CheckpointStateOutputStream org.apache.flink.runtime.state.CheckpointStateOutputStream used on opensource projects

Example 1 with CheckpointStateOutputStream

use of org.apache.flink.runtime.state.CheckpointStateOutputStream in project flink by apache.

the class HeapSnapshotStrategy method asyncSnapshot.

@Override
public SnapshotResultSupplier<KeyedStateHandle> asyncSnapshot(HeapSnapshotResources<K> syncPartResource, long checkpointId, long timestamp, @Nonnull CheckpointStreamFactory streamFactory, @Nonnull CheckpointOptions checkpointOptions) {
    List<StateMetaInfoSnapshot> metaInfoSnapshots = syncPartResource.getMetaInfoSnapshots();
    if (metaInfoSnapshots.isEmpty()) {
        return snapshotCloseableRegistry -> SnapshotResult.empty();
    }
    final KeyedBackendSerializationProxy<K> serializationProxy = new KeyedBackendSerializationProxy<>(// get a serialized form already at state registration time in the future
    syncPartResource.getKeySerializer(), metaInfoSnapshots, !Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, keyGroupCompressionDecorator));
    final SupplierWithException<CheckpointStreamWithResultProvider, Exception> checkpointStreamSupplier = localRecoveryConfig.isLocalRecoveryEnabled() && !checkpointOptions.getCheckpointType().isSavepoint() ? () -> createDuplicatingStream(checkpointId, CheckpointedStateScope.EXCLUSIVE, streamFactory, localRecoveryConfig.getLocalStateDirectoryProvider().orElseThrow(LocalRecoveryConfig.localRecoveryNotEnabled())) : () -> createSimpleStream(CheckpointedStateScope.EXCLUSIVE, streamFactory);
    return (snapshotCloseableRegistry) -> {
        final Map<StateUID, Integer> stateNamesToId = syncPartResource.getStateNamesToId();
        final Map<StateUID, StateSnapshot> cowStateStableSnapshots = syncPartResource.getCowStateStableSnapshots();
        final CheckpointStreamWithResultProvider streamWithResultProvider = checkpointStreamSupplier.get();
        snapshotCloseableRegistry.registerCloseable(streamWithResultProvider);
        final CheckpointStateOutputStream localStream = streamWithResultProvider.getCheckpointOutputStream();
        final DataOutputViewStreamWrapper outView = new DataOutputViewStreamWrapper(localStream);
        serializationProxy.write(outView);
        final long[] keyGroupRangeOffsets = new long[keyGroupRange.getNumberOfKeyGroups()];
        for (int keyGroupPos = 0; keyGroupPos < keyGroupRange.getNumberOfKeyGroups(); ++keyGroupPos) {
            int keyGroupId = keyGroupRange.getKeyGroupId(keyGroupPos);
            keyGroupRangeOffsets[keyGroupPos] = localStream.getPos();
            outView.writeInt(keyGroupId);
            for (Map.Entry<StateUID, StateSnapshot> stateSnapshot : cowStateStableSnapshots.entrySet()) {
                StateSnapshot.StateKeyGroupWriter partitionedSnapshot = stateSnapshot.getValue().getKeyGroupWriter();
                try (OutputStream kgCompressionOut = keyGroupCompressionDecorator.decorateWithCompression(localStream)) {
                    DataOutputViewStreamWrapper kgCompressionView = new DataOutputViewStreamWrapper(kgCompressionOut);
                    kgCompressionView.writeShort(stateNamesToId.get(stateSnapshot.getKey()));
                    partitionedSnapshot.writeStateInKeyGroup(kgCompressionView, keyGroupId);
                }
            // this will just close the outer compression stream
            }
        }
        if (snapshotCloseableRegistry.unregisterCloseable(streamWithResultProvider)) {
            KeyGroupRangeOffsets kgOffs = new KeyGroupRangeOffsets(keyGroupRange, keyGroupRangeOffsets);
            SnapshotResult<StreamStateHandle> result = streamWithResultProvider.closeAndFinalizeCheckpointStreamResult();
            return toKeyedStateHandleSnapshotResult(result, kgOffs, KeyGroupsStateHandle::new);
        } else {
            throw new IOException("Stream already unregistered.");
        }
    };
}

Example 2 with CheckpointStateOutputStream

use of org.apache.flink.runtime.state.CheckpointStateOutputStream in project flink by apache.

the class MemoryCheckpointOutputStreamTest method testOversizedState.

@Test
public void testOversizedState() throws Exception {
    HashMap<String, Integer> state = new HashMap<>();
    state.put("hey there", 2);
    state.put("the crazy brown fox stumbles over a sentence that does not contain every letter", 77);
    CheckpointStateOutputStream outStream = new MemoryCheckpointOutputStream(10);
    ObjectOutputStream oos = new ObjectOutputStream(outStream);
    oos.writeObject(state);
    oos.flush();
    try {
        outStream.closeAndGetHandle();
        fail("this should cause an exception");
    } catch (IOException e) {
    // that's what we expect
    }
}

Example 3 with CheckpointStateOutputStream

use of org.apache.flink.runtime.state.CheckpointStateOutputStream in project flink by apache.

the class StateSnapshotContextSynchronousImplTest method testStreamClosingWhenClosing.

/**
 * Tests that closing the StateSnapshotContextSynchronousImpl will also close the associated
 * output streams.
 */
@Test
public void testStreamClosingWhenClosing() throws Exception {
    long checkpointId = 42L;
    long checkpointTimestamp = 1L;
    CheckpointStateOutputStream outputStream1 = mock(CheckpointStateOutputStream.class);
    CheckpointStateOutputStream outputStream2 = mock(CheckpointStateOutputStream.class);
    CheckpointStreamFactory streamFactory = mock(CheckpointStreamFactory.class);
    when(streamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE)).thenReturn(outputStream1, outputStream2);
    InsightCloseableRegistry closableRegistry = new InsightCloseableRegistry();
    KeyGroupRange keyGroupRange = new KeyGroupRange(0, 2);
    StateSnapshotContextSynchronousImpl context = new StateSnapshotContextSynchronousImpl(checkpointId, checkpointTimestamp, streamFactory, keyGroupRange, closableRegistry);
    // creating the output streams
    context.getRawKeyedOperatorStateOutput();
    context.getRawOperatorStateOutput();
    verify(streamFactory, times(2)).createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    assertEquals(2, closableRegistry.size());
    assertTrue(closableRegistry.contains(outputStream1));
    assertTrue(closableRegistry.contains(outputStream2));
    context.getKeyedStateStreamFuture().run();
    context.getOperatorStateStreamFuture().run();
    verify(outputStream1).closeAndGetHandle();
    verify(outputStream2).closeAndGetHandle();
    assertEquals(0, closableRegistry.size());
}

Example 4 with CheckpointStateOutputStream

use of org.apache.flink.runtime.state.CheckpointStateOutputStream in project flink by apache.

the class StateSnapshotContextSynchronousImplTest method testStreamClosingExceptionally.

@Test
public void testStreamClosingExceptionally() throws Exception {
    long checkpointId = 42L;
    long checkpointTimestamp = 1L;
    CheckpointStateOutputStream outputStream1 = mock(CheckpointStateOutputStream.class);
    CheckpointStateOutputStream outputStream2 = mock(CheckpointStateOutputStream.class);
    CheckpointStreamFactory streamFactory = mock(CheckpointStreamFactory.class);
    when(streamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE)).thenReturn(outputStream1, outputStream2);
    InsightCloseableRegistry closableRegistry = new InsightCloseableRegistry();
    KeyGroupRange keyGroupRange = new KeyGroupRange(0, 2);
    StateSnapshotContextSynchronousImpl context = new StateSnapshotContextSynchronousImpl(checkpointId, checkpointTimestamp, streamFactory, keyGroupRange, closableRegistry);
    // creating the output streams
    context.getRawKeyedOperatorStateOutput();
    context.getRawOperatorStateOutput();
    verify(streamFactory, times(2)).createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    assertEquals(2, closableRegistry.size());
    assertTrue(closableRegistry.contains(outputStream1));
    assertTrue(closableRegistry.contains(outputStream2));
    context.closeExceptionally();
    verify(outputStream1).close();
    verify(outputStream2).close();
    assertEquals(0, closableRegistry.size());
}

Example 5 with CheckpointStateOutputStream

use of org.apache.flink.runtime.state.CheckpointStateOutputStream in project flink by apache.

the class FsCheckpointStateOutputStreamTest method testMixedBelowAndAboveThreshold.

@Test
public void testMixedBelowAndAboveThreshold() throws Exception {
    final byte[] state1 = new byte[1274673];
    final byte[] state2 = new byte[1];
    final byte[] state3 = new byte[0];
    final byte[] state4 = new byte[177];
    final Random rnd = new Random();
    rnd.nextBytes(state1);
    rnd.nextBytes(state2);
    rnd.nextBytes(state3);
    rnd.nextBytes(state4);
    final File directory = tempDir.newFolder();
    final Path basePath = Path.fromLocalFile(directory);
    final Supplier<CheckpointStateOutputStream> factory = () -> new FsCheckpointStateOutputStream(basePath, FileSystem.getLocalFileSystem(), 1024, 15, relativePaths);
    CheckpointStateOutputStream stream1 = factory.get();
    CheckpointStateOutputStream stream2 = factory.get();
    CheckpointStateOutputStream stream3 = factory.get();
    stream1.write(state1);
    stream2.write(state2);
    stream3.write(state3);
    FileStateHandle handle1 = (FileStateHandle) stream1.closeAndGetHandle();
    ByteStreamStateHandle handle2 = (ByteStreamStateHandle) stream2.closeAndGetHandle();
    ByteStreamStateHandle handle3 = (ByteStreamStateHandle) stream3.closeAndGetHandle();
    // use with try-with-resources
    StreamStateHandle handle4;
    try (CheckpointStateOutputStream stream4 = factory.get()) {
        stream4.write(state4);
        handle4 = stream4.closeAndGetHandle();
    }
    // close before accessing handle
    CheckpointStateOutputStream stream5 = factory.get();
    stream5.write(state4);
    stream5.close();
    try {
        stream5.closeAndGetHandle();
        fail();
    } catch (IOException e) {
    // uh-huh
    }
    validateBytesInStream(handle1.openInputStream(), state1);
    handle1.discardState();
    assertFalse(isDirectoryEmpty(directory));
    ensureLocalFileDeleted(handle1.getFilePath());
    validateBytesInStream(handle2.openInputStream(), state2);
    handle2.discardState();
    assertFalse(isDirectoryEmpty(directory));
    // nothing was written to the stream, so it will return nothing
    assertNull(handle3);
    assertFalse(isDirectoryEmpty(directory));
    validateBytesInStream(handle4.openInputStream(), state4);
    handle4.discardState();
    assertTrue(isDirectoryEmpty(directory));
}