void wakeUp(boolean taskOnly) {
        // Synchronize to make sure the wake up only works for the current invocation of runOnce().
        lock.lock();
        try {
            wakeUpUnsafe(taskOnly);
        } finally {
            lock.unlock();
        }
    }
@Test
    public void testWakeup() throws InterruptedException {
        final int numSplits = 3;
        final int numRecordsPerSplit = 10_000;
        final int wakeupRecordsInterval = 10;
        final int numTotalRecords = numRecordsPerSplit * numSplits;

        FutureCompletingBlockingQueue<RecordsWithSplitIds<int[]>> elementQueue =
                new FutureCompletingBlockingQueue<>(1);
        SplitFetcher<int[], MockSourceSplit> fetcher =
                new SplitFetcher<>(
                        0,
                        elementQueue,
                        MockSplitReader.newBuilder()
                                .setNumRecordsPerSplitPerFetch(2)
                                .setBlockingFetch(true)
                                .build(),
                        ExceptionUtils::rethrow,
                        () -> {},
                        (ignore) -> {},
                        false);

        // Prepare the splits.
        List<MockSourceSplit> splits = new ArrayList<>();
        for (int i = 0; i < numSplits; i++) {
            splits.add(new MockSourceSplit(i, 0, numRecordsPerSplit));
            int base = i * numRecordsPerSplit;
            for (int j = base; j < base + numRecordsPerSplit; j++) {
                splits.get(splits.size() - 1).addRecord(j);
            }
        }
        // Add splits to the fetcher.
        fetcher.addSplits(splits);

        // A thread drives the fetcher.
        Thread fetcherThread = new Thread(fetcher, "FetcherThread");

        SortedSet<Integer> recordsRead = Collections.synchronizedSortedSet(new TreeSet<>());

        // A thread waking up the split fetcher frequently.
        AtomicInteger wakeupTimes = new AtomicInteger(0);
        AtomicBoolean stop = new AtomicBoolean(false);
        Thread wakeUpCaller =
                new Thread("Wakeup Caller") {
                    @Override
                    public void run() {
                        int lastWakeup = 0;
                        while (recordsRead.size() < numTotalRecords && !stop.get()) {
                            int numRecordsRead = recordsRead.size();
                            if (numRecordsRead >= lastWakeup + wakeupRecordsInterval) {
                                fetcher.wakeUp(false);
                                wakeupTimes.incrementAndGet();
                                lastWakeup = numRecordsRead;
                            }
                        }
                    }
                };

        try {
            fetcherThread.start();
            wakeUpCaller.start();

            while (recordsRead.size() < numSplits * numRecordsPerSplit) {
                final RecordsWithSplitIds<int[]> nextBatch = elementQueue.take();
                while (nextBatch.nextSplit() != null) {
                    int[] arr;
                    while ((arr = nextBatch.nextRecordFromSplit()) != null) {
                        assertThat(recordsRead.add(arr[0])).isTrue();
                    }
                }
            }

            assertThat(recordsRead).hasSize(numTotalRecords);
            assertThat(recordsRead.first()).isEqualTo(0);
            assertThat(recordsRead.last()).isEqualTo(numTotalRecords - 1);
            assertThat(wakeupTimes.get()).isGreaterThan(0);
        } finally {
            stop.set(true);
            fetcher.shutdown();
            fetcherThread.join();
            wakeUpCaller.join();
        }
    }