boolean reportSubtaskStats(JobVertexID jobVertexId, SubtaskStateStats subtask) {
        TaskStateStats taskStateStats = taskStats.get(jobVertexId);

        if (taskStateStats != null && taskStateStats.reportSubtaskStats(subtask)) {
            if (subtask.isCompleted()) {
                currentNumAcknowledgedSubtasks++;
                latestAcknowledgedSubtask = subtask;
            }

            currentCheckpointedSize += subtask.getCheckpointedSize();
            currentStateSize += subtask.getStateSize();

            long processedData = subtask.getProcessedData();
            if (processedData > 0) {
                currentProcessedData += processedData;
            }

            long persistedData = subtask.getPersistedData();
            if (persistedData > 0) {
                currentPersistedData += persistedData;
            }

            unalignedCheckpoint |= subtask.getUnalignedCheckpoint();

            return true;
        } else {
            return false;
        }
    }
@Test
    void testReportSubtaskStats() {
        long checkpointId = Integer.MAX_VALUE + 1222L;
        long triggerTimestamp = Integer.MAX_VALUE - 1239L;
        CheckpointProperties props =
                CheckpointProperties.forCheckpoint(
                        CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION);
        TaskStateStats task1 = new TaskStateStats(new JobVertexID(), 3);
        TaskStateStats task2 = new TaskStateStats(new JobVertexID(), 4);
        int totalSubtaskCount = task1.getNumberOfSubtasks() + task2.getNumberOfSubtasks();

        HashMap<JobVertexID, TaskStateStats> taskStats = new HashMap<>();
        taskStats.put(task1.getJobVertexId(), task1);
        taskStats.put(task2.getJobVertexId(), task2);

        PendingCheckpointStats pending =
                new PendingCheckpointStats(
                        checkpointId, triggerTimestamp, props, totalSubtaskCount, taskStats);

        // Check initial state
        assertThat(pending.getCheckpointId()).isEqualTo(checkpointId);
        assertThat(pending.getTriggerTimestamp()).isEqualTo(triggerTimestamp);
        assertThat(pending.getProperties()).isEqualTo(props);
        assertThat(pending.getStatus()).isEqualTo(CheckpointStatsStatus.IN_PROGRESS);
        assertThat(pending.getNumberOfAcknowledgedSubtasks()).isZero();
        assertThat(pending.getStateSize()).isZero();
        assertThat(pending.getNumberOfSubtasks()).isEqualTo(totalSubtaskCount);
        assertThat(pending.getLatestAcknowledgedSubtaskStats()).isNull();
        assertThat(pending.getLatestAckTimestamp()).isEqualTo(-1);
        assertThat(pending.getEndToEndDuration()).isEqualTo(-1);
        assertThat(pending.getTaskStateStats(task1.getJobVertexId())).isEqualTo(task1);
        assertThat(pending.getTaskStateStats(task2.getJobVertexId())).isEqualTo(task2);
        assertThat(pending.getTaskStateStats(new JobVertexID())).isNull();

        // Report subtasks and check getters
        assertThat(pending.reportSubtaskStats(new JobVertexID(), createSubtaskStats(0, false)))
                .isFalse();

        long stateSize = 0;

        // Report 1st task
        for (int i = 0; i < task1.getNumberOfSubtasks(); i++) {
            SubtaskStateStats subtask = createSubtaskStats(i, false);
            stateSize += subtask.getStateSize();

            pending.reportSubtaskStats(task1.getJobVertexId(), subtask);

            assertThat(pending.isUnalignedCheckpoint()).isFalse();
            assertThat(pending.getLatestAcknowledgedSubtaskStats()).isEqualTo(subtask);
            assertThat(pending.getLatestAckTimestamp()).isEqualTo(subtask.getAckTimestamp());
            assertThat(pending.getEndToEndDuration())
                    .isEqualTo(subtask.getAckTimestamp() - triggerTimestamp);
            assertThat(pending.getStateSize()).isEqualTo(stateSize);
        }

        // Don't allow overwrite
        assertThat(pending.reportSubtaskStats(task1.getJobVertexId(), task1.getSubtaskStats()[0]))
                .isFalse();

        // Report 2nd task
        for (int i = 0; i < task2.getNumberOfSubtasks(); i++) {
            SubtaskStateStats subtask = createSubtaskStats(i, true);
            stateSize += subtask.getStateSize();

            pending.reportSubtaskStats(task2.getJobVertexId(), subtask);

            assertThat(pending.isUnalignedCheckpoint()).isTrue();
            assertThat(pending.getLatestAcknowledgedSubtaskStats()).isEqualTo(subtask);
            assertThat(pending.getLatestAckTimestamp()).isEqualTo(subtask.getAckTimestamp());
            assertThat(pending.getEndToEndDuration())
                    .isEqualTo(subtask.getAckTimestamp() - triggerTimestamp);
            assertThat(pending.getStateSize()).isEqualTo(stateSize);
        }

        assertThat(task1.getNumberOfAcknowledgedSubtasks()).isEqualTo(task1.getNumberOfSubtasks());
        assertThat(task2.getNumberOfAcknowledgedSubtasks()).isEqualTo(task2.getNumberOfSubtasks());
    }