Examples with MockSourceSplit org.apache.flink.api.connector.source.mocks.MockSourceSplit used on opensource projects

Example 26 with MockSourceSplit

use of org.apache.flink.api.connector.source.mocks.MockSourceSplit in project flink by apache.

the class SplitFetcherTest method testWakeup.

@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) -> {
    });
    // 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) {
                    assertTrue(recordsRead.add(arr[0]));
                }
            }
        }
        assertEquals(numTotalRecords, recordsRead.size());
        assertEquals(0, (int) recordsRead.first());
        assertEquals(numTotalRecords - 1, (int) recordsRead.last());
        assertTrue(wakeupTimes.get() > 0);
    } finally {
        stop.set(true);
        fetcher.shutdown();
        fetcherThread.join();
        wakeUpCaller.join();
    }
}

Example 27 with MockSourceSplit

use of org.apache.flink.api.connector.source.mocks.MockSourceSplit in project flink by apache.

the class HybridSourceReaderTest method testReaderRecovery.

@Test
public void testReaderRecovery() throws Exception {
    TestingReaderContext readerContext = new TestingReaderContext();
    TestingReaderOutput<Integer> readerOutput = new TestingReaderOutput<>();
    MockBaseSource source = new MockBaseSource(1, 1, Boundedness.BOUNDED);
    HybridSourceReader<Integer> reader = new HybridSourceReader<>(readerContext);
    reader.start();
    assertAndClearSourceReaderFinishedEvent(readerContext, -1);
    reader.handleSourceEvents(new SwitchSourceEvent(0, source, false));
    MockSourceSplit mockSplit = new MockSourceSplit(0, 0, 2147483647);
    SwitchedSources switchedSources = new SwitchedSources();
    switchedSources.put(0, source);
    HybridSourceSplit hybridSplit = HybridSourceSplit.wrapSplit(mockSplit, 0, switchedSources);
    reader.addSplits(Collections.singletonList(hybridSplit));
    List<HybridSourceSplit> snapshot = reader.snapshotState(0);
    Assert.assertThat(snapshot, Matchers.contains(hybridSplit));
    // reader recovery
    readerContext.clearSentEvents();
    reader = new HybridSourceReader<>(readerContext);
    reader.addSplits(snapshot);
    Assert.assertNull(currentReader(reader));
    reader.start();
    Assert.assertNull(currentReader(reader));
    assertAndClearSourceReaderFinishedEvent(readerContext, -1);
    reader.handleSourceEvents(new SwitchSourceEvent(0, source, false));
    Assert.assertNotNull(currentReader(reader));
    Assert.assertThat(reader.snapshotState(1), Matchers.contains(hybridSplit));
    reader.close();
}

Example 28 with MockSourceSplit

use of org.apache.flink.api.connector.source.mocks.MockSourceSplit in project flink by apache.

the class HybridSourceReaderTest method testDefaultMethodDelegation.

@Test
public void testDefaultMethodDelegation() throws Exception {
    TestingReaderContext readerContext = new TestingReaderContext();
    TestingReaderOutput<Integer> readerOutput = new TestingReaderOutput<>();
    MockBaseSource source = new MockBaseSource(1, 1, Boundedness.BOUNDED) {

        @Override
        public SourceReader<Integer, MockSourceSplit> createReader(SourceReaderContext readerContext) {
            return Mockito.spy(super.createReader(readerContext));
        }
    };
    HybridSourceReader<Integer> reader = new HybridSourceReader<>(readerContext);
    reader.start();
    assertAndClearSourceReaderFinishedEvent(readerContext, -1);
    reader.handleSourceEvents(new SwitchSourceEvent(0, source, false));
    SourceReader<Integer, MockSourceSplit> underlyingReader = currentReader(reader);
    reader.notifyCheckpointComplete(1);
    Mockito.verify(underlyingReader).notifyCheckpointComplete(1);
    reader.notifyCheckpointAborted(1);
    Mockito.verify(underlyingReader).notifyCheckpointAborted(1);
    reader.close();
}

Example 29 with MockSourceSplit

use of org.apache.flink.api.connector.source.mocks.MockSourceSplit in project flink by apache.

the class MockBaseSource method createEnumerator.

@Override
public SplitEnumerator<MockSourceSplit, List<MockSourceSplit>> createEnumerator(SplitEnumeratorContext<MockSourceSplit> enumContext) {
    List<MockSourceSplit> splits = new ArrayList<>();
    for (int i = 0; i < numSplits; i++) {
        int endIndex = boundedness == Boundedness.BOUNDED ? numRecordsPerSplit : Integer.MAX_VALUE;
        MockSourceSplit split = new MockSourceSplit(i, 0, endIndex);
        for (int j = 0; j < numRecordsPerSplit; j++) {
            split.addRecord(startingValue + i * numRecordsPerSplit + j);
        }
        splits.add(split);
    }
    return new MockSplitEnumerator(splits, enumContext);
}

Example 30 with MockSourceSplit

use of org.apache.flink.api.connector.source.mocks.MockSourceSplit in project flink by apache.

the class MockSplitReader method getRecords.

private RecordsBySplits<int[]> getRecords() {
    final RecordsBySplits.Builder<int[]> records = new RecordsBySplits.Builder<>();
    // after this locked section, the thread might be interrupted
    synchronized (wakeupLock) {
        if (wokenUp) {
            wokenUp = false;
            return records.build();
        }
        threadInBlocking = Thread.currentThread();
    }
    try {
        Iterator<Map.Entry<String, MockSourceSplit>> iterator = splits.entrySet().iterator();
        while (iterator.hasNext()) {
            Map.Entry<String, MockSourceSplit> entry = iterator.next();
            MockSourceSplit split = entry.getValue();
            boolean hasRecords = false;
            for (int i = 0; i < numRecordsPerSplitPerFetch && !split.isFinished(); i++) {
                // This call may throw InterruptedException.
                int[] record = split.getNext(blockingFetch);
                if (record != null) {
                    records.add(entry.getKey(), record);
                    hasRecords = true;
                }
            }
            if (split.isFinished()) {
                if (!separatedFinishedRecord) {
                    records.addFinishedSplit(entry.getKey());
                    iterator.remove();
                } else if (!hasRecords) {
                    records.addFinishedSplit(entry.getKey());
                    iterator.remove();
                    break;
                }
            }
        }
    } catch (InterruptedException ie) {
        // Catch the exception and return the records that are already read.
        if (!blockingFetch) {
            throw new RuntimeException("Caught unexpected interrupted exception.");
        }
    } finally {
        // after this locked section, the thread may not be interrupted any more
        synchronized (wakeupLock) {
            wokenUp = false;
            // noinspection ResultOfMethodCallIgnored
            Thread.interrupted();
            threadInBlocking = null;
        }
    }
    return records.build();
}