Examples with RecordsBySplits org.apache.flink.connector.base.source.reader.RecordsBySplits used on opensource projects

Example 1 with RecordsBySplits

use of org.apache.flink.connector.base.source.reader.RecordsBySplits 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();
}