/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.ignite.internal.pagememory.persistence.checkpoint;

import static java.util.concurrent.CompletableFuture.completedFuture;
import static java.util.concurrent.TimeUnit.SECONDS;
import static java.util.stream.Collectors.toList;
import static org.apache.ignite.internal.pagememory.persistence.checkpoint.CheckpointDirtyPages.EMPTY;
import static org.apache.ignite.internal.pagememory.persistence.checkpoint.CheckpointState.FINISHED;
import static org.apache.ignite.internal.pagememory.persistence.checkpoint.CheckpointState.LOCK_RELEASED;
import static org.apache.ignite.internal.pagememory.persistence.checkpoint.CheckpointState.LOCK_TAKEN;
import static org.apache.ignite.internal.pagememory.persistence.checkpoint.CheckpointState.PAGES_SNAPSHOT_TAKEN;
import static org.apache.ignite.internal.pagememory.persistence.checkpoint.CheckpointState.PAGES_SORTED;
import static org.apache.ignite.internal.pagememory.persistence.checkpoint.CheckpointTestUtils.newReadWriteLock;
import static org.apache.ignite.internal.pagememory.persistence.checkpoint.CheckpointTestUtils.toListDirtyPageIds;
import static org.apache.ignite.internal.pagememory.persistence.checkpoint.CheckpointWorkflowTest.TestCheckpointListener.AFTER_CHECKPOINT_END;
import static org.apache.ignite.internal.pagememory.persistence.checkpoint.CheckpointWorkflowTest.TestCheckpointListener.BEFORE_CHECKPOINT_BEGIN;
import static org.apache.ignite.internal.pagememory.persistence.checkpoint.CheckpointWorkflowTest.TestCheckpointListener.ON_CHECKPOINT_BEGIN;
import static org.apache.ignite.internal.pagememory.persistence.checkpoint.CheckpointWorkflowTest.TestCheckpointListener.ON_MARK_CHECKPOINT_BEGIN;
import static org.apache.ignite.internal.testframework.IgniteTestUtils.await;
import static org.apache.ignite.internal.testframework.IgniteTestUtils.runAsync;
import static org.apache.ignite.internal.util.FastTimestamps.coarseCurrentTimeMillis;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.empty;
import static org.hamcrest.Matchers.equalTo;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.isNull;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.stream.IntStream;
import org.apache.ignite.internal.logger.IgniteLogger;
import org.apache.ignite.internal.logger.Loggers;
import org.apache.ignite.internal.pagememory.DataRegion;
import org.apache.ignite.internal.pagememory.FullPageId;
import org.apache.ignite.internal.pagememory.persistence.PersistentPageMemory;
import org.apache.ignite.internal.pagememory.persistence.checkpoint.CheckpointDirtyPages.CheckpointDirtyPagesView;
import org.apache.ignite.internal.pagememory.util.PageIdUtils;
import org.apache.ignite.lang.IgniteInternalCheckedException;
import org.jetbrains.annotations.Nullable;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;

/**
 * For {@link CheckpointWorkflow} testing.
 */
public class CheckpointWorkflowTest {
    private final IgniteLogger log = Loggers.forClass(CheckpointWorkflowTest.class);

    @Nullable
    private CheckpointWorkflow workflow;

    @AfterEach
    void tearDown() {
        if (workflow != null) {
            workflow.stop();
        }
    }

    @Test
    void testListeners() {
        PersistentPageMemory pageMemory0 = mock(PersistentPageMemory.class);
        PersistentPageMemory pageMemory1 = mock(PersistentPageMemory.class);
        PersistentPageMemory pageMemory2 = mock(PersistentPageMemory.class);

        DataRegion<PersistentPageMemory> dataRegion0 = () -> pageMemory0;
        DataRegion<PersistentPageMemory> dataRegion1 = () -> pageMemory1;
        DataRegion<PersistentPageMemory> dataRegion2 = () -> pageMemory2;

        workflow = new CheckpointWorkflow(
                "test",
                newReadWriteLock(log),
                List.of(dataRegion0, dataRegion1),
                1
        );

        workflow.start();

        CheckpointListener listener0 = mock(CheckpointListener.class);
        CheckpointListener listener1 = mock(CheckpointListener.class);
        CheckpointListener listener2 = mock(CheckpointListener.class);
        CheckpointListener listener3 = mock(CheckpointListener.class);
        CheckpointListener listener4 = mock(CheckpointListener.class);

        workflow.addCheckpointListener(listener0, dataRegion0);
        workflow.addCheckpointListener(listener1, dataRegion0);
        workflow.addCheckpointListener(listener2, dataRegion0);
        workflow.addCheckpointListener(listener3, dataRegion1);
        workflow.addCheckpointListener(listener4, null);

        assertThat(
                Set.copyOf(workflow.collectCheckpointListeners(List.of(dataRegion0))),
                equalTo(Set.of(listener0, listener1, listener2, listener4))
        );

        assertThat(
                Set.copyOf(workflow.collectCheckpointListeners(List.of(dataRegion1))),
                equalTo(Set.of(listener3, listener4))
        );

        assertThat(
                Set.copyOf(workflow.collectCheckpointListeners(List.of(dataRegion2))),
                equalTo(Set.of(listener4))
        );

        assertThat(
                Set.copyOf(workflow.collectCheckpointListeners(List.of(dataRegion0, dataRegion1))),
                equalTo(Set.of(listener0, listener1, listener2, listener3, listener4))
        );

        // Checks remove listener.

        workflow.removeCheckpointListener(listener0);
        workflow.removeCheckpointListener(listener1);
        workflow.removeCheckpointListener(listener3);

        assertThat(
                Set.copyOf(workflow.collectCheckpointListeners(List.of(dataRegion0))),
                equalTo(Set.of(listener2, listener4))
        );

        assertThat(
                Set.copyOf(workflow.collectCheckpointListeners(List.of(dataRegion1))),
                equalTo(Set.of(listener4))
        );

        assertThat(
                Set.copyOf(workflow.collectCheckpointListeners(List.of(dataRegion2))),
                equalTo(Set.of(listener4))
        );

        // Checks empty listeners after stop.

        workflow.stop();

        assertThat(workflow.collectCheckpointListeners(List.of(dataRegion0)), empty());
        assertThat(workflow.collectCheckpointListeners(List.of(dataRegion1)), empty());
        assertThat(workflow.collectCheckpointListeners(List.of(dataRegion2)), empty());
    }

    @Test
    void testMarkCheckpointBegin() throws Exception {
        CheckpointReadWriteLock readWriteLock = newReadWriteLock(log);

        List<FullPageId> dirtyPages = List.of(of(0, 0, 1), of(0, 0, 2), of(0, 0, 3));

        PersistentPageMemory pageMemory = newPageMemory(dirtyPages);

        DataRegion<PersistentPageMemory> dataRegion = () -> pageMemory;

        workflow = new CheckpointWorkflow(
                "test",
                readWriteLock,
                List.of(dataRegion),
                1
        );

        workflow.start();

        CheckpointProgressImpl progressImpl = mock(CheckpointProgressImpl.class);

        ArgumentCaptor<CheckpointState> checkpointStateArgumentCaptor = ArgumentCaptor.forClass(CheckpointState.class);

        ArgumentCaptor<Integer> pagesCountArgumentCaptor = ArgumentCaptor.forClass(Integer.class);

        doNothing().when(progressImpl).transitTo(checkpointStateArgumentCaptor.capture());

        doNothing().when(progressImpl).currentCheckpointPagesCount(pagesCountArgumentCaptor.capture());

        when(progressImpl.futureFor(PAGES_SORTED)).thenReturn(completedFuture(null));

        UUID checkpointId = UUID.randomUUID();

        when(progressImpl.id()).thenReturn(checkpointId);

        List<String> events = new ArrayList<>();

        CheckpointMetricsTracker tracker = mock(CheckpointMetricsTracker.class);

        Runnable onReleaseWriteLock = mock(Runnable.class);

        workflow.addCheckpointListener(new TestCheckpointListener(events) {
            /** {@inheritDoc} */
            @Override
            public void beforeCheckpointBegin(CheckpointProgress progress, @Nullable Executor exec) throws IgniteInternalCheckedException {
                super.beforeCheckpointBegin(progress, exec);

                assertNull(exec);

                assertSame(progressImpl, progress);

                assertEquals(readWriteLock.getReadHoldCount(), 1);

                assertThat(checkpointStateArgumentCaptor.getAllValues(), empty());

                verify(tracker, never()).onWriteLockWaitStart();
                verify(tracker, never()).onMarkCheckpointBeginStart();
                verify(tracker, never()).onMarkCheckpointBeginEnd();
                verify(tracker, never()).onWriteLockRelease();
                verify(tracker, never()).onSplitAndSortCheckpointPagesStart();
                verify(tracker, never()).onSplitAndSortCheckpointPagesEnd();

                verify(progressImpl, never()).pagesToWrite(any(CheckpointDirtyPages.class));
                verify(progressImpl, never()).initCounters(anyInt());

                verify(onReleaseWriteLock, never()).run();
            }

            /** {@inheritDoc} */
            @Override
            public void onMarkCheckpointBegin(CheckpointProgress progress, @Nullable Executor exec) throws IgniteInternalCheckedException {
                super.onMarkCheckpointBegin(progress, exec);

                assertNull(exec);

                assertSame(progressImpl, progress);

                assertTrue(readWriteLock.isWriteLockHeldByCurrentThread());

                assertThat(checkpointStateArgumentCaptor.getAllValues(), equalTo(List.of(LOCK_TAKEN)));

                verify(tracker, times(1)).onWriteLockWaitStart();
                verify(tracker, times(1)).onMarkCheckpointBeginStart();
                verify(tracker, never()).onMarkCheckpointBeginEnd();
                verify(tracker, never()).onWriteLockRelease();
                verify(tracker, never()).onSplitAndSortCheckpointPagesStart();
                verify(tracker, never()).onSplitAndSortCheckpointPagesEnd();

                verify(progressImpl, never()).pagesToWrite(any(CheckpointDirtyPages.class));
                verify(progressImpl, never()).initCounters(anyInt());

                verify(onReleaseWriteLock, never()).run();
            }

            /** {@inheritDoc} */
            @Override
            public void onCheckpointBegin(CheckpointProgress progress) throws IgniteInternalCheckedException {
                super.onCheckpointBegin(progress);

                assertSame(progressImpl, progress);

                assertEquals(0, readWriteLock.getReadHoldCount());

                assertFalse(readWriteLock.isWriteLockHeldByCurrentThread());

                assertThat(checkpointStateArgumentCaptor.getAllValues(), equalTo(List.of(LOCK_TAKEN, PAGES_SNAPSHOT_TAKEN, LOCK_RELEASED)));

                assertThat(pagesCountArgumentCaptor.getAllValues(), equalTo(List.of(3)));

                verify(tracker, times(1)).onWriteLockWaitStart();
                verify(tracker, times(1)).onMarkCheckpointBeginStart();
                verify(tracker, times(1)).onMarkCheckpointBeginEnd();
                verify(tracker, times(1)).onWriteLockRelease();
                verify(tracker, never()).onSplitAndSortCheckpointPagesStart();
                verify(tracker, never()).onSplitAndSortCheckpointPagesEnd();

                verify(progressImpl, never()).pagesToWrite(any(CheckpointDirtyPages.class));
                verify(progressImpl, never()).initCounters(anyInt());

                verify(onReleaseWriteLock, times(1)).run();
            }
        }, dataRegion);

        Checkpoint checkpoint = workflow.markCheckpointBegin(
                coarseCurrentTimeMillis(),
                progressImpl,
                tracker,
                () -> {},
                onReleaseWriteLock
        );

        verify(tracker, times(1)).onWriteLockWaitStart();
        verify(tracker, times(1)).onMarkCheckpointBeginStart();
        verify(tracker, times(1)).onMarkCheckpointBeginEnd();
        verify(tracker, times(1)).onWriteLockRelease();
        verify(tracker, times(1)).onSplitAndSortCheckpointPagesStart();
        verify(tracker, times(1)).onSplitAndSortCheckpointPagesEnd();

        verify(progressImpl, times(1)).pagesToWrite(any(CheckpointDirtyPages.class));
        verify(progressImpl, times(1)).initCounters(anyInt());

        CheckpointDirtyPagesView dirtyPagesView = checkpoint.dirtyPages.nextPartitionView(null);

        assertThat(toListDirtyPageIds(dirtyPagesView), equalTo(dirtyPages));
        assertThat(dirtyPagesView.pageMemory(), equalTo(dataRegion.pageMemory()));

        assertThat(
                events,
                equalTo(List.of(BEFORE_CHECKPOINT_BEGIN, ON_MARK_CHECKPOINT_BEGIN, ON_CHECKPOINT_BEGIN))
        );

        assertThat(
                checkpointStateArgumentCaptor.getAllValues(),
                equalTo(List.of(LOCK_TAKEN, PAGES_SNAPSHOT_TAKEN, LOCK_RELEASED, PAGES_SORTED))
        );

        verify(onReleaseWriteLock, times(1)).run();
    }

    @Test
    void testMarkCheckpointEnd() throws Exception {
        CheckpointReadWriteLock readWriteLock = newReadWriteLock(log);

        PersistentPageMemory pageMemory = mock(PersistentPageMemory.class);

        DataRegion<PersistentPageMemory> dataRegion = () -> pageMemory;

        workflow = new CheckpointWorkflow(
                "test",
                readWriteLock,
                List.of(dataRegion),
                1
        );

        workflow.start();

        List<String> events = new ArrayList<>();

        ArgumentCaptor<CheckpointState> checkpointStateArgumentCaptor = ArgumentCaptor.forClass(CheckpointState.class);

        CheckpointProgressImpl progressImpl = mock(CheckpointProgressImpl.class);

        doNothing().when(progressImpl).transitTo(checkpointStateArgumentCaptor.capture());

        UUID checkpointId = UUID.randomUUID();

        when(progressImpl.id()).thenReturn(checkpointId);

        TestCheckpointListener checkpointListener = new TestCheckpointListener(events) {
            /** {@inheritDoc} */
            @Override
            public void afterCheckpointEnd(CheckpointProgress progress) throws IgniteInternalCheckedException {
                super.afterCheckpointEnd(progress);

                assertSame(progressImpl, progress);

                assertFalse(readWriteLock.isWriteLockHeldByCurrentThread());

                assertEquals(0, readWriteLock.getReadHoldCount());

                assertThat(checkpointStateArgumentCaptor.getAllValues(), empty());

                verify(progressImpl, times(1)).pagesToWrite(isNull());
                verify(progressImpl, times(1)).clearCounters();
            }
        };

        workflow.addCheckpointListener(checkpointListener, dataRegion);

        workflow.markCheckpointEnd(new Checkpoint(
                new CheckpointDirtyPages(List.of(createCheckpointDirtyPages(pageMemory, of(0, 0, 0)))),
                progressImpl
        ));

        assertThat(checkpointStateArgumentCaptor.getAllValues(), equalTo(List.of(FINISHED)));

        assertThat(events, equalTo(List.of(AFTER_CHECKPOINT_END)));

        verify(progressImpl, times(1)).clearCounters();

        verify(pageMemory, times(1)).finishCheckpoint();

        verify(progressImpl, times(1)).pagesToWrite(isNull());

        verify(progressImpl, times(1)).clearCounters();

        // Checks with empty dirty pages.

        workflow.removeCheckpointListener(checkpointListener);

        assertDoesNotThrow(() -> workflow.markCheckpointEnd(new Checkpoint(EMPTY, progressImpl)));
    }

    @Test
    void testCreateAndSortCheckpointDirtyPages() throws Exception {
        DataRegionDirtyPages<Collection<FullPageId>> dataRegionDirtyPages0 = createDataRegionDirtyPages(
                mock(PersistentPageMemory.class),
                of(10, 10, 2), of(10, 10, 1), of(10, 10, 0),
                of(10, 5, 100), of(10, 5, 99),
                of(10, 1, 50), of(10, 1, 51), of(10, 1, 99)
        );

        DataRegionDirtyPages<Collection<FullPageId>> dataRegionDirtyPages1 = createDataRegionDirtyPages(
                mock(PersistentPageMemory.class),
                of(77, 5, 100), of(77, 5, 99),
                of(88, 1, 51), of(88, 1, 50), of(88, 1, 99),
                of(66, 33, 0), of(66, 33, 1), of(66, 33, 2)
        );

        workflow = new CheckpointWorkflow(
                "test",
                newReadWriteLock(log),
                List.of(),
                1
        );

        workflow.start();

        CheckpointDirtyPages sortCheckpointDirtyPages = workflow.createAndSortCheckpointDirtyPages(
                new DataRegionsDirtyPages(List.of(dataRegionDirtyPages0, dataRegionDirtyPages1))
        );

        CheckpointDirtyPagesView dirtyPagesView = sortCheckpointDirtyPages.nextPartitionView(null);

        assertThat(toListDirtyPageIds(dirtyPagesView), equalTo(List.of(of(10, 1, 50), of(10, 1, 51), of(10, 1, 99))));
        assertThat(dirtyPagesView.pageMemory(), equalTo(dataRegionDirtyPages0.pageMemory));

        dirtyPagesView = sortCheckpointDirtyPages.nextPartitionView(dirtyPagesView);

        assertThat(toListDirtyPageIds(dirtyPagesView), equalTo(List.of(of(10, 5, 99), of(10, 5, 100))));
        assertThat(dirtyPagesView.pageMemory(), equalTo(dataRegionDirtyPages0.pageMemory));

        dirtyPagesView = sortCheckpointDirtyPages.nextPartitionView(dirtyPagesView);

        assertThat(toListDirtyPageIds(dirtyPagesView), equalTo(List.of(of(10, 10, 0), of(10, 10, 1), of(10, 10, 2))));
        assertThat(dirtyPagesView.pageMemory(), equalTo(dataRegionDirtyPages0.pageMemory));

        dirtyPagesView = sortCheckpointDirtyPages.nextPartitionView(dirtyPagesView);

        assertThat(toListDirtyPageIds(dirtyPagesView), equalTo(List.of(of(66, 33, 0), of(66, 33, 1), of(66, 33, 2))));
        assertThat(dirtyPagesView.pageMemory(), equalTo(dataRegionDirtyPages1.pageMemory));

        dirtyPagesView = sortCheckpointDirtyPages.nextPartitionView(dirtyPagesView);

        assertThat(toListDirtyPageIds(dirtyPagesView), equalTo(List.of(of(77, 5, 99), of(77, 5, 100))));
        assertThat(dirtyPagesView.pageMemory(), equalTo(dataRegionDirtyPages1.pageMemory));

        dirtyPagesView = sortCheckpointDirtyPages.nextPartitionView(dirtyPagesView);

        assertThat(toListDirtyPageIds(dirtyPagesView), equalTo(List.of(of(88, 1, 50), of(88, 1, 51), of(88, 1, 99))));
        assertThat(dirtyPagesView.pageMemory(), equalTo(dataRegionDirtyPages1.pageMemory));
    }

    @Test
    void testParallelSortDirtyPages() throws Exception {
        int count = CheckpointWorkflow.PARALLEL_SORT_THRESHOLD + 10;

        FullPageId[] dirtyPages0 = IntStream.range(0, count).mapToObj(i -> of(0, 0, count - i)).toArray(FullPageId[]::new);
        FullPageId[] dirtyPages1 = IntStream.range(0, count).mapToObj(i -> of(1, 1, i)).toArray(FullPageId[]::new);

        workflow = new CheckpointWorkflow(
                "test",
                newReadWriteLock(log),
                List.of(),
                1
        );

        workflow.start();

        CheckpointDirtyPages sortCheckpointDirtyPages = workflow.createAndSortCheckpointDirtyPages(new DataRegionsDirtyPages(List.of(
                createDataRegionDirtyPages(mock(PersistentPageMemory.class), dirtyPages1),
                createDataRegionDirtyPages(mock(PersistentPageMemory.class), dirtyPages0)
        )));

        CheckpointDirtyPagesView dirtyPagesView = sortCheckpointDirtyPages.nextPartitionView(null);

        assertThat(toListDirtyPageIds(dirtyPagesView), equalTo(List.of(dirtyPages1)));

        dirtyPagesView = sortCheckpointDirtyPages.nextPartitionView(dirtyPagesView);

        assertThat(
                toListDirtyPageIds(dirtyPagesView),
                equalTo(IntStream.range(0, count).mapToObj(i -> dirtyPages0[count - i - 1]).collect(toList()))
        );
    }

    @Test
    void testAwaitPendingTasksOfListenerCallback() {
        workflow = new CheckpointWorkflow(
                "test",
                newReadWriteLock(log),
                List.of(),
                2
        );

        workflow.start();

        CompletableFuture<?> startTaskBeforeCheckpointBeginFuture = new CompletableFuture<>();
        CompletableFuture<?> finishTaskBeforeCheckpointBeginFuture = new CompletableFuture<>();

        CompletableFuture<?> startTaskOnMarkCheckpointBeginFuture = new CompletableFuture<>();
        CompletableFuture<?> finishTaskOnMarkCheckpointBeginFuture = new CompletableFuture<>();

        workflow.addCheckpointListener(new CheckpointListener() {
            /** {@inheritDoc} */
            @Override
            public void beforeCheckpointBegin(CheckpointProgress progress, @Nullable Executor executor) {
                assertNotNull(executor);

                executor.execute(() -> {
                    startTaskBeforeCheckpointBeginFuture.complete(null);

                    await(finishTaskBeforeCheckpointBeginFuture, 1, SECONDS);
                });
            }

            /** {@inheritDoc} */
            @Override
            public void onMarkCheckpointBegin(CheckpointProgress progress, @Nullable Executor executor) {
                assertNotNull(executor);

                executor.execute(() -> {
                    startTaskOnMarkCheckpointBeginFuture.complete(null);

                    await(finishTaskOnMarkCheckpointBeginFuture, 1, SECONDS);
                });
            }
        }, null);

        Runnable updateHeartbeat = mock(Runnable.class);

        CompletableFuture<Checkpoint> markCheckpointBeginFuture = runAsync(() -> workflow.markCheckpointBegin(
                coarseCurrentTimeMillis(),
                mock(CheckpointProgressImpl.class),
                mock(CheckpointMetricsTracker.class),
                updateHeartbeat,
                () -> {}
        ));

        await(startTaskBeforeCheckpointBeginFuture, 1, SECONDS);

        assertFalse(markCheckpointBeginFuture.isDone());
        assertFalse(startTaskOnMarkCheckpointBeginFuture.isDone());
        verify(updateHeartbeat, times(1)).run();

        finishTaskBeforeCheckpointBeginFuture.complete(null);

        await(startTaskOnMarkCheckpointBeginFuture, 1, SECONDS);

        assertFalse(markCheckpointBeginFuture.isDone());
        verify(updateHeartbeat, times(3)).run();

        finishTaskOnMarkCheckpointBeginFuture.complete(null);

        await(markCheckpointBeginFuture, 1, SECONDS);

        verify(updateHeartbeat, times(5)).run();
    }

    private static PersistentPageMemory newPageMemory(Collection<FullPageId> pageIds) {
        PersistentPageMemory mock = mock(PersistentPageMemory.class);

        when(mock.beginCheckpoint(any(CompletableFuture.class))).thenReturn(pageIds);

        return mock;
    }

    private static DataRegionDirtyPages<FullPageId[]> createCheckpointDirtyPages(
            PersistentPageMemory pageMemory,
            FullPageId... pageIds
    ) {
        return new DataRegionDirtyPages<>(pageMemory, pageIds);
    }

    private static DataRegionDirtyPages<Collection<FullPageId>> createDataRegionDirtyPages(
            PersistentPageMemory pageMemory,
            FullPageId... pageIds
    ) {
        return new DataRegionDirtyPages<>(pageMemory, List.of(pageIds));
    }

    private static FullPageId of(int grpId, int partId, int pageIdx) {
        return new FullPageId(PageIdUtils.pageId(partId, (byte) 0, pageIdx), grpId);
    }

    /**
     * Test listener implementation that simply collects events.
     */
    static class TestCheckpointListener implements CheckpointListener {
        static final String ON_MARK_CHECKPOINT_BEGIN = "onMarkCheckpointBegin";

        static final String ON_CHECKPOINT_BEGIN = "onCheckpointBegin";

        static final String BEFORE_CHECKPOINT_BEGIN = "beforeCheckpointBegin";

        static final String AFTER_CHECKPOINT_END = "afterCheckpointEnd";

        final List<String> events;

        /**
         * Constructor.
         *
         * @param events For recording events.
         */
        TestCheckpointListener(List<String> events) {
            this.events = events;
        }

        /** {@inheritDoc} */
        @Override
        public void onMarkCheckpointBegin(CheckpointProgress progress, @Nullable Executor executor) throws IgniteInternalCheckedException {
            events.add(ON_MARK_CHECKPOINT_BEGIN);
        }

        /** {@inheritDoc} */
        @Override
        public void onCheckpointBegin(CheckpointProgress progress) throws IgniteInternalCheckedException {
            events.add(ON_CHECKPOINT_BEGIN);
        }

        /** {@inheritDoc} */
        @Override
        public void beforeCheckpointBegin(CheckpointProgress progress, @Nullable Executor executor) throws IgniteInternalCheckedException {
            events.add(BEFORE_CHECKPOINT_BEGIN);
        }

        /** {@inheritDoc} */
        @Override
        public void afterCheckpointEnd(CheckpointProgress progress) throws IgniteInternalCheckedException {
            events.add(AFTER_CHECKPOINT_END);
        }
    }
}
