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

Example 1 with RecordsWithSplitIds

use of org.apache.flink.connector.base.source.reader.RecordsWithSplitIds in project flink by apache.

the class SplitFetcherManagerTest method testExceptionPropagation.

// the final modifier is important so that '@SafeVarargs' is accepted on Java 8
@SuppressWarnings("FinalPrivateMethod")
@SafeVarargs
private final void testExceptionPropagation(final RecordsWithSplitIds<Integer>... fetchesBeforeError) throws Exception {
    final IOException testingException = new IOException("test");
    final FutureCompletingBlockingQueue<RecordsWithSplitIds<Integer>> queue = new FutureCompletingBlockingQueue<>(10);
    final AwaitingReader<Integer, TestingSourceSplit> reader = new AwaitingReader<>(testingException, fetchesBeforeError);
    final SplitFetcherManager<Integer, TestingSourceSplit> fetcher = createFetcher("testSplit", queue, reader);
    reader.awaitAllRecordsReturned();
    drainQueue(queue);
    assertFalse(queue.getAvailabilityFuture().isDone());
    reader.triggerThrowException();
    // await the error propagation
    queue.getAvailabilityFuture().get();
    try {
        fetcher.checkErrors();
        fail("expected exception");
    } catch (Exception e) {
        assertSame(testingException, e.getCause().getCause());
    } finally {
        fetcher.close(20_000L);
    }
}

Example 2 with RecordsWithSplitIds

use of org.apache.flink.connector.base.source.reader.RecordsWithSplitIds in project flink by apache.

the class KafkaSource method createReader.

@VisibleForTesting
SourceReader<OUT, KafkaPartitionSplit> createReader(SourceReaderContext readerContext, Consumer<Collection<String>> splitFinishedHook) throws Exception {
    FutureCompletingBlockingQueue<RecordsWithSplitIds<ConsumerRecord<byte[], byte[]>>> elementsQueue = new FutureCompletingBlockingQueue<>();
    deserializationSchema.open(new DeserializationSchema.InitializationContext() {

        @Override
        public MetricGroup getMetricGroup() {
            return readerContext.metricGroup().addGroup("deserializer");
        }

        @Override
        public UserCodeClassLoader getUserCodeClassLoader() {
            return readerContext.getUserCodeClassLoader();
        }
    });
    final KafkaSourceReaderMetrics kafkaSourceReaderMetrics = new KafkaSourceReaderMetrics(readerContext.metricGroup());
    Supplier<KafkaPartitionSplitReader> splitReaderSupplier = () -> new KafkaPartitionSplitReader(props, readerContext, kafkaSourceReaderMetrics);
    KafkaRecordEmitter<OUT> recordEmitter = new KafkaRecordEmitter<>(deserializationSchema);
    return new KafkaSourceReader<>(elementsQueue, new KafkaSourceFetcherManager(elementsQueue, splitReaderSupplier::get, splitFinishedHook), recordEmitter, toConfiguration(props), readerContext, kafkaSourceReaderMetrics);
}

Example 3 with RecordsWithSplitIds

use of org.apache.flink.connector.base.source.reader.RecordsWithSplitIds in project flink by apache.

the class PulsarSourceReaderFactory method create.

@SuppressWarnings("java:S2095")
public static <OUT> SourceReader<OUT, PulsarPartitionSplit> create(SourceReaderContext readerContext, PulsarDeserializationSchema<OUT> deserializationSchema, SourceConfiguration sourceConfiguration) {
    PulsarClient pulsarClient = createClient(sourceConfiguration);
    PulsarAdmin pulsarAdmin = createAdmin(sourceConfiguration);
    // Create a message queue with the predefined source option.
    int queueCapacity = sourceConfiguration.getMessageQueueCapacity();
    FutureCompletingBlockingQueue<RecordsWithSplitIds<PulsarMessage<OUT>>> elementsQueue = new FutureCompletingBlockingQueue<>(queueCapacity);
    // Create different pulsar source reader by subscription type.
    SubscriptionType subscriptionType = sourceConfiguration.getSubscriptionType();
    if (subscriptionType == SubscriptionType.Failover || subscriptionType == SubscriptionType.Exclusive) {
        // Create a ordered split reader supplier.
        Supplier<PulsarOrderedPartitionSplitReader<OUT>> splitReaderSupplier = () -> new PulsarOrderedPartitionSplitReader<>(pulsarClient, pulsarAdmin, sourceConfiguration, deserializationSchema);
        return new PulsarOrderedSourceReader<>(elementsQueue, splitReaderSupplier, readerContext, sourceConfiguration, pulsarClient, pulsarAdmin);
    } else if (subscriptionType == SubscriptionType.Shared || subscriptionType == SubscriptionType.Key_Shared) {
        TransactionCoordinatorClient coordinatorClient = ((PulsarClientImpl) pulsarClient).getTcClient();
        if (coordinatorClient == null && !sourceConfiguration.isEnableAutoAcknowledgeMessage()) {
            throw new IllegalStateException("Transaction is required but didn't enabled");
        }
        Supplier<PulsarUnorderedPartitionSplitReader<OUT>> splitReaderSupplier = () -> new PulsarUnorderedPartitionSplitReader<>(pulsarClient, pulsarAdmin, sourceConfiguration, deserializationSchema, coordinatorClient);
        return new PulsarUnorderedSourceReader<>(elementsQueue, splitReaderSupplier, readerContext, sourceConfiguration, pulsarClient, pulsarAdmin, coordinatorClient);
    } else {
        throw new UnsupportedOperationException("This subscription type is not " + subscriptionType + " supported currently.");
    }
}

Example 4 with RecordsWithSplitIds

use of org.apache.flink.connector.base.source.reader.RecordsWithSplitIds 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 5 with RecordsWithSplitIds

use of org.apache.flink.connector.base.source.reader.RecordsWithSplitIds in project flink by apache.

the class MockBaseSource method createReader.

@Override
public SourceReader<Integer, MockSourceSplit> createReader(SourceReaderContext readerContext) {
    FutureCompletingBlockingQueue<RecordsWithSplitIds<int[]>> elementsQueue = new FutureCompletingBlockingQueue<>();
    Configuration config = new Configuration();
    config.setInteger(SourceReaderOptions.ELEMENT_QUEUE_CAPACITY, 1);
    config.setLong(SourceReaderOptions.SOURCE_READER_CLOSE_TIMEOUT, 30000L);
    MockSplitReader.Builder builder = MockSplitReader.newBuilder().setNumRecordsPerSplitPerFetch(2).setBlockingFetch(true);
    return new MockSourceReader(elementsQueue, builder::build, config, readerContext);
}