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

Example 11 with CheckpointStateOutputStream

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

the class MemoryCheckpointOutputStreamTest method testStateStream.

@Test
public void testStateStream() 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(MemoryStateBackend.DEFAULT_MAX_STATE_SIZE);
    ObjectOutputStream oos = new ObjectOutputStream(outStream);
    oos.writeObject(state);
    oos.flush();
    StreamStateHandle handle = outStream.closeAndGetHandle();
    assertNotNull(handle);
    try (ObjectInputStream ois = new ObjectInputStream(handle.openInputStream())) {
        assertEquals(state, ois.readObject());
        assertTrue(ois.available() <= 0);
    }
}

Example 12 with CheckpointStateOutputStream

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

the class MemoryCheckpointStorageAccessTest method testTaskOwnedStateStream.

@Test
public void testTaskOwnedStateStream() throws Exception {
    final List<String> state = Arrays.asList("Flopsy", "Mopsy", "Cotton Tail", "Peter");
    final MemoryBackendCheckpointStorageAccess storage = new MemoryBackendCheckpointStorageAccess(new JobID(), null, null, DEFAULT_MAX_STATE_SIZE);
    StreamStateHandle stateHandle;
    try (CheckpointStateOutputStream stream = storage.createTaskOwnedStateStream()) {
        assertTrue(stream instanceof MemoryCheckpointOutputStream);
        new ObjectOutputStream(stream).writeObject(state);
        stateHandle = stream.closeAndGetHandle();
    }
    try (ObjectInputStream in = new ObjectInputStream(stateHandle.openInputStream())) {
        assertEquals(state, in.readObject());
    }
}

Example 13 with CheckpointStateOutputStream

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

the class FsCheckpointStorageAccessTest method testTaskOwnedStateStream.

@Test
public void testTaskOwnedStateStream() throws Exception {
    final List<String> state = Arrays.asList("Flopsy", "Mopsy", "Cotton Tail", "Peter");
    // we chose a small size threshold here to force the state to disk
    final FsCheckpointStorageAccess storage = new FsCheckpointStorageAccess(Path.fromLocalFile(tmp.newFolder()), null, new JobID(), 10, WRITE_BUFFER_SIZE);
    final StreamStateHandle stateHandle;
    try (CheckpointStateOutputStream stream = storage.createTaskOwnedStateStream()) {
        assertTrue(stream instanceof FsCheckpointStateOutputStream);
        new ObjectOutputStream(stream).writeObject(state);
        stateHandle = stream.closeAndGetHandle();
    }
    // the state must have gone to disk
    FileStateHandle fileStateHandle = (FileStateHandle) stateHandle;
    // check that the state is in the correct directory
    String parentDirName = fileStateHandle.getFilePath().getParent().getName();
    assertEquals(FsCheckpointStorageAccess.CHECKPOINT_TASK_OWNED_STATE_DIR, parentDirName);
    // validate the contents
    try (ObjectInputStream in = new ObjectInputStream(stateHandle.openInputStream())) {
        assertEquals(state, in.readObject());
    }
}

Example 14 with CheckpointStateOutputStream

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

the class RocksDBStateUploader method uploadLocalFileToCheckpointFs.

private StreamStateHandle uploadLocalFileToCheckpointFs(Path filePath, CheckpointStreamFactory checkpointStreamFactory, CheckpointedStateScope stateScope, CloseableRegistry closeableRegistry) throws IOException {
    InputStream inputStream = null;
    CheckpointStateOutputStream outputStream = null;
    try {
        final byte[] buffer = new byte[READ_BUFFER_SIZE];
        inputStream = Files.newInputStream(filePath);
        closeableRegistry.registerCloseable(inputStream);
        outputStream = checkpointStreamFactory.createCheckpointStateOutputStream(stateScope);
        closeableRegistry.registerCloseable(outputStream);
        while (true) {
            int numBytes = inputStream.read(buffer);
            if (numBytes == -1) {
                break;
            }
            outputStream.write(buffer, 0, numBytes);
        }
        StreamStateHandle result = null;
        if (closeableRegistry.unregisterCloseable(outputStream)) {
            result = outputStream.closeAndGetHandle();
            outputStream = null;
        }
        return result;
    } finally {
        if (closeableRegistry.unregisterCloseable(inputStream)) {
            IOUtils.closeQuietly(inputStream);
        }
        if (closeableRegistry.unregisterCloseable(outputStream)) {
            IOUtils.closeQuietly(outputStream);
        }
    }
}

Example 15 with CheckpointStateOutputStream

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

the class FsCheckpointStateOutputStreamTest method testCleanupWhenClosingStream.

/**
 * Tests that the underlying stream file is deleted upon calling close.
 */
@Test
public void testCleanupWhenClosingStream() throws IOException {
    final FileSystem fs = mock(FileSystem.class);
    final FSDataOutputStream outputStream = mock(FSDataOutputStream.class);
    final ArgumentCaptor<Path> pathCaptor = ArgumentCaptor.forClass(Path.class);
    when(fs.create(pathCaptor.capture(), any(FileSystem.WriteMode.class))).thenReturn(outputStream);
    CheckpointStateOutputStream stream = new FsCheckpointStreamFactory.FsCheckpointStateOutputStream(Path.fromLocalFile(tempDir.newFolder()), fs, 4, 0, relativePaths);
    // this should create the underlying file stream
    stream.write(new byte[] { 1, 2, 3, 4, 5 });
    verify(fs).create(any(Path.class), any(FileSystem.WriteMode.class));
    stream.close();
    verify(fs).delete(eq(pathCaptor.getValue()), anyBoolean());
}