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

Example 6 with CheckpointStreamFactory

use of org.apache.flink.runtime.state.CheckpointStreamFactory 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;
    CheckpointStreamFactory.CheckpointStateOutputStream outputStream1 = mock(CheckpointStreamFactory.CheckpointStateOutputStream.class);
    CheckpointStreamFactory.CheckpointStateOutputStream outputStream2 = mock(CheckpointStreamFactory.CheckpointStateOutputStream.class);
    CheckpointStreamFactory streamFactory = mock(CheckpointStreamFactory.class);
    when(streamFactory.createCheckpointStateOutputStream(eq(checkpointId), eq(checkpointTimestamp))).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(eq(checkpointId), eq(checkpointTimestamp));
    assertEquals(2, closableRegistry.size());
    assertTrue(closableRegistry.contains(outputStream1));
    assertTrue(closableRegistry.contains(outputStream2));
    context.close();
    verify(outputStream1).close();
    verify(outputStream2).close();
    assertEquals(0, closableRegistry.size());
}

Example 7 with CheckpointStreamFactory

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

the class SavepointV0Serializer method convertKeyedBackendState.

/**
	 * This is public so that we can use it when restoring a legacy snapshot
	 * in {@code AbstractStreamOperatorTestHarness}.
	 */
public static KeyGroupsStateHandle convertKeyedBackendState(HashMap<String, KvStateSnapshot<?, ?, ?, ?>> oldKeyedState, int parallelInstanceIdx, long checkpointID) throws Exception {
    if (null != oldKeyedState) {
        CheckpointStreamFactory checkpointStreamFactory = new MemCheckpointStreamFactory(MAX_SIZE);
        CheckpointStreamFactory.CheckpointStateOutputStream keyedStateOut = checkpointStreamFactory.createCheckpointStateOutputStream(checkpointID, 0L);
        try {
            final long offset = keyedStateOut.getPos();
            InstantiationUtil.serializeObject(keyedStateOut, oldKeyedState);
            StreamStateHandle streamStateHandle = keyedStateOut.closeAndGetHandle();
            // makes IOUtils.closeQuietly(...) ignore this
            keyedStateOut = null;
            if (null != streamStateHandle) {
                KeyGroupRangeOffsets keyGroupRangeOffsets = new KeyGroupRangeOffsets(parallelInstanceIdx, parallelInstanceIdx, new long[] { offset });
                return new MigrationKeyGroupStateHandle(keyGroupRangeOffsets, streamStateHandle);
            }
        } finally {
            IOUtils.closeQuietly(keyedStateOut);
        }
    }
    return null;
}

Example 8 with CheckpointStreamFactory

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

the class KeyedOneInputStreamOperatorTestHarness method snapshotLegacy.

/**
	 *
	 */
@Override
public StreamStateHandle snapshotLegacy(long checkpointId, long timestamp) throws Exception {
    // simply use an in-memory handle
    MemoryStateBackend backend = new MemoryStateBackend();
    CheckpointStreamFactory streamFactory = backend.createStreamFactory(new JobID(), "test_op");
    CheckpointStreamFactory.CheckpointStateOutputStream outStream = streamFactory.createCheckpointStateOutputStream(checkpointId, timestamp);
    if (operator instanceof StreamCheckpointedOperator) {
        ((StreamCheckpointedOperator) operator).snapshotState(outStream, checkpointId, timestamp);
    }
    if (keyedStateBackend != null) {
        RunnableFuture<KeyGroupsStateHandle> keyedSnapshotRunnable = keyedStateBackend.snapshot(checkpointId, timestamp, streamFactory, CheckpointOptions.forFullCheckpoint());
        if (!keyedSnapshotRunnable.isDone()) {
            Thread runner = new Thread(keyedSnapshotRunnable);
            runner.start();
        }
        outStream.write(1);
        ObjectOutputStream oos = new ObjectOutputStream(outStream);
        oos.writeObject(keyedSnapshotRunnable.get());
        oos.flush();
    } else {
        outStream.write(0);
    }
    return outStream.closeAndGetHandle();
}

Example 9 with CheckpointStreamFactory

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

the class FileStateBackendTest method testStateOutputStream.

@Test
public void testStateOutputStream() {
    URI basePath = randomHdfsFileUri();
    try {
        FsStateBackend backend = CommonTestUtils.createCopySerializable(new FsStateBackend(basePath, 15));
        JobID jobId = new JobID();
        CheckpointStreamFactory streamFactory = backend.createStreamFactory(jobId, "test_op");
        // we know how FsCheckpointStreamFactory is implemented so we know where it
        // will store checkpoints
        Path checkpointPath = new Path(new Path(basePath), jobId.toString());
        byte[] state1 = new byte[1274673];
        byte[] state2 = new byte[1];
        byte[] state3 = new byte[0];
        byte[] state4 = new byte[177];
        Random rnd = new Random();
        rnd.nextBytes(state1);
        rnd.nextBytes(state2);
        rnd.nextBytes(state3);
        rnd.nextBytes(state4);
        long checkpointId = 97231523452L;
        CheckpointStreamFactory.CheckpointStateOutputStream stream1 = streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
        CheckpointStreamFactory.CheckpointStateOutputStream stream2 = streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
        CheckpointStreamFactory.CheckpointStateOutputStream stream3 = streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
        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
        FileStateHandle handle4;
        try (CheckpointStreamFactory.CheckpointStateOutputStream stream4 = streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis())) {
            stream4.write(state4);
            handle4 = (FileStateHandle) stream4.closeAndGetHandle();
        }
        // close before accessing handle
        CheckpointStreamFactory.CheckpointStateOutputStream stream5 = streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
        stream5.write(state4);
        stream5.close();
        try {
            stream5.closeAndGetHandle();
            fail();
        } catch (IOException e) {
        // uh-huh
        }
        validateBytesInStream(handle1.openInputStream(), state1);
        handle1.discardState();
        assertFalse(isDirectoryEmpty(checkpointPath));
        ensureFileDeleted(handle1.getFilePath());
        validateBytesInStream(handle2.openInputStream(), state2);
        handle2.discardState();
        // stream 3 has zero bytes, so it should not return anything
        assertNull(handle3);
        validateBytesInStream(handle4.openInputStream(), state4);
        handle4.discardState();
        assertTrue(isDirectoryEmpty(checkpointPath));
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}

Example 10 with CheckpointStreamFactory

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

the class RocksDBAsyncSnapshotTest method testCleanupOfSnapshotsInFailureCase.

/**
	 * Test that the snapshot files are cleaned up in case of a failure during the snapshot
	 * procedure.
	 */
@Test
public void testCleanupOfSnapshotsInFailureCase() throws Exception {
    long checkpointId = 1L;
    long timestamp = 42L;
    Environment env = new DummyEnvironment("test task", 1, 0);
    CheckpointStreamFactory.CheckpointStateOutputStream outputStream = mock(CheckpointStreamFactory.CheckpointStateOutputStream.class);
    CheckpointStreamFactory checkpointStreamFactory = mock(CheckpointStreamFactory.class);
    AbstractStateBackend stateBackend = mock(AbstractStateBackend.class);
    final IOException testException = new IOException("Test exception");
    doReturn(checkpointStreamFactory).when(stateBackend).createStreamFactory(any(JobID.class), anyString());
    doThrow(testException).when(outputStream).write(anyInt());
    doReturn(outputStream).when(checkpointStreamFactory).createCheckpointStateOutputStream(eq(checkpointId), eq(timestamp));
    RocksDBStateBackend backend = new RocksDBStateBackend(stateBackend);
    backend.setDbStoragePath("file:///tmp/foobar");
    AbstractKeyedStateBackend<Void> keyedStateBackend = backend.createKeyedStateBackend(env, new JobID(), "test operator", VoidSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), null);
    // register a state so that the state backend has to checkpoint something
    keyedStateBackend.getPartitionedState("namespace", StringSerializer.INSTANCE, new ValueStateDescriptor<>("foobar", String.class));
    RunnableFuture<KeyGroupsStateHandle> snapshotFuture = keyedStateBackend.snapshot(checkpointId, timestamp, checkpointStreamFactory, CheckpointOptions.forFullCheckpoint());
    try {
        FutureUtil.runIfNotDoneAndGet(snapshotFuture);
        fail("Expected an exception to be thrown here.");
    } catch (ExecutionException e) {
        Assert.assertEquals(testException, e.getCause());
    }
    verify(outputStream).close();
}