@Override
    public FileMergingCheckpointStateOutputStream createCheckpointStateOutputStream(
            SubtaskKey subtaskKey, long checkpointId, CheckpointedStateScope scope) {

        return new FileMergingCheckpointStateOutputStream(
                writeBufferSize,
                new FileMergingCheckpointStateOutputStream.FileMergingSnapshotManagerProxy() {
                    PhysicalFile physicalFile;
                    LogicalFile logicalFile;

                    @Override
                    public Tuple2<FSDataOutputStream, Path> providePhysicalFile()
                            throws IOException {
                        physicalFile =
                                getOrCreatePhysicalFileForCheckpoint(
                                        subtaskKey, checkpointId, scope);
                        return new Tuple2<>(
                                physicalFile.getOutputStream(), physicalFile.getFilePath());
                    }

                    @Override
                    public SegmentFileStateHandle closeStreamAndCreateStateHandle(
                            Path filePath, long startPos, long stateSize) throws IOException {
                        if (physicalFile == null) {
                            return null;
                        } else {
                            // deal with logical file
                            logicalFile =
                                    createLogicalFile(
                                            physicalFile, startPos, stateSize, subtaskKey);
                            logicalFile.advanceLastCheckpointId(checkpointId);

                            // track the logical file
                            synchronized (lock) {
                                uploadedStates
                                        .computeIfAbsent(checkpointId, key -> new HashSet<>())
                                        .add(logicalFile);
                            }

                            // deal with physicalFile file
                            returnPhysicalFileForNextReuse(subtaskKey, checkpointId, physicalFile);

                            return new SegmentFileStateHandle(
                                    physicalFile.getFilePath(),
                                    startPos,
                                    stateSize,
                                    scope,
                                    logicalFile.getFileId());
                        }
                    }

                    @Override
                    public void closeStreamExceptionally() throws IOException {
                        if (physicalFile != null) {
                            if (logicalFile != null) {
                                discardSingleLogicalFile(logicalFile, checkpointId);
                            } else {
                                // The physical file should be closed anyway. This is because the
                                // last segmented write on this file is likely to have failed, and
                                // we want to prevent further reusing of this file.
                                physicalFile.close();
                                physicalFile.deleteIfNecessary();
                            }
                        }
                    }
                });
    }
@Test
    public void testReusedFileWriting() throws Exception {
        long checkpointId = 1;
        int streamNum = 10;
        int perStreamWriteNum = 128;

        // write random bytes and then read them from the file
        byte[] bytes = new byte[streamNum * perStreamWriteNum];
        Random rd = new Random();
        rd.nextBytes(bytes);
        int byteIndex = 0;

        SegmentFileStateHandle[] handles = new SegmentFileStateHandle[streamNum];
        try (FileMergingSnapshotManager fmsm = createFileMergingSnapshotManager(checkpointBaseDir);
                CloseableRegistry closeableRegistry = new CloseableRegistry()) {

            // repeatedly get-write-close streams
            for (int i = 0; i < streamNum; i++) {
                FileMergingCheckpointStateOutputStream stream =
                        fmsm.createCheckpointStateOutputStream(
                                subtaskKey1, checkpointId, CheckpointedStateScope.EXCLUSIVE);
                try {
                    closeableRegistry.registerCloseable(stream);
                    for (int j = 0; j < perStreamWriteNum; j++) {
                        stream.write(bytes[byteIndex++]);
                    }
                    handles[i] = stream.closeAndGetHandle();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }

            // assert the streams writes to the same file correctly
            byteIndex = 0;
            Path filePath = null;
            for (SegmentFileStateHandle handle : handles) {
                // check file path
                Path thisFilePath = handle.getFilePath();
                assertThat(filePath == null || filePath.equals(thisFilePath)).isTrue();
                filePath = thisFilePath;
                // check file content
                FSDataInputStream is = handle.openInputStream();

                closeableRegistry.registerCloseable(is);
                int readValue;

                while ((readValue = is.read()) != -1) {
                    assertThat((byte) readValue).isEqualTo(bytes[byteIndex++]);
                }
            }
        }
    }