org.apache.flink.connector.base.source.reader.fetcher.SplitFetcherManager Maven / Gradle / Ivy
/*
* 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.flink.connector.base.source.reader.fetcher;
import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.api.connector.source.SourceSplit;
import org.apache.flink.connector.base.source.reader.RecordsWithSplitIds;
import org.apache.flink.connector.base.source.reader.SourceReaderBase;
import org.apache.flink.connector.base.source.reader.splitreader.SplitReader;
import org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import java.util.function.Supplier;
/**
* A class responsible for starting the {@link SplitFetcher} and manage the life cycles of them.
* This class works with the {@link SourceReaderBase}.
*
* The split fetcher manager could be used to support different threading models by implementing
* the {@link #addSplits(List)} method differently. For example, a single thread split fetcher
* manager would only start a single fetcher and assign all the splits to it. A one-thread-per-split
* fetcher may spawn a new thread every time a new split is assigned.
*/
public abstract class SplitFetcherManager {
private static final Logger LOG = LoggerFactory.getLogger(SplitFetcherManager.class);
private final Consumer errorHandler;
/** An atomic integer to generate monotonically increasing fetcher ids. */
private final AtomicInteger fetcherIdGenerator;
/** A supplier to provide split readers. */
private final Supplier> splitReaderFactory;
/** Uncaught exception in the split fetchers. */
private final AtomicReference uncaughtFetcherException;
/** The element queue that the split fetchers will put elements into. */
private final FutureCompletingBlockingQueue> elementsQueue;
/** A map keeping track of all the split fetchers. */
protected final Map> fetchers;
/**
* An executor service with two threads. One for the fetcher and one for the future completing
* thread.
*/
private final ExecutorService executors;
/** Indicating the split fetcher manager has closed or not. */
private volatile boolean closed;
/**
* Hook for handling finished splits in {@link SplitFetcher}, usually used for testing split
* finishing behavior of {@link SplitFetcher} and {@link SplitReader}.
*/
private final Consumer> splitFinishedHook;
/**
* Create a split fetcher manager.
*
* @param elementsQueue the queue that split readers will put elements into.
* @param splitReaderFactory a supplier that could be used to create split readers.
*/
public SplitFetcherManager(
FutureCompletingBlockingQueue> elementsQueue,
Supplier> splitReaderFactory) {
this(elementsQueue, splitReaderFactory, (ignore) -> {});
}
/**
* Create a split fetcher manager.
*
* @param elementsQueue the queue that split readers will put elements into.
* @param splitReaderFactory a supplier that could be used to create split readers.
* @param splitFinishedHook Hook for handling finished splits in split fetchers.
*/
@VisibleForTesting
public SplitFetcherManager(
FutureCompletingBlockingQueue> elementsQueue,
Supplier> splitReaderFactory,
Consumer> splitFinishedHook) {
this.elementsQueue = elementsQueue;
this.errorHandler =
new Consumer() {
@Override
public void accept(Throwable t) {
LOG.error("Received uncaught exception.", t);
if (!uncaughtFetcherException.compareAndSet(null, t)) {
// Add the exception to the exception list.
uncaughtFetcherException.get().addSuppressed(t);
}
// Wake up the main thread to let it know the exception.
elementsQueue.notifyAvailable();
}
};
this.splitReaderFactory = splitReaderFactory;
this.splitFinishedHook = splitFinishedHook;
this.uncaughtFetcherException = new AtomicReference<>(null);
this.fetcherIdGenerator = new AtomicInteger(0);
this.fetchers = new ConcurrentHashMap<>();
// Create the executor with a thread factory that fails the source reader if one of
// the fetcher thread exits abnormally.
final String taskThreadName = Thread.currentThread().getName();
this.executors =
Executors.newCachedThreadPool(
r -> new Thread(r, "Source Data Fetcher for " + taskThreadName));
this.closed = false;
}
public abstract void addSplits(List splitsToAdd);
protected void startFetcher(SplitFetcher fetcher) {
executors.submit(fetcher);
}
/**
* Synchronize method to ensure no fetcher is created after the split fetcher manager has
* closed.
*
* @return the created split fetcher.
* @throws IllegalStateException if the split fetcher manager has closed.
*/
protected synchronized SplitFetcher createSplitFetcher() {
if (closed) {
throw new IllegalStateException("The split fetcher manager has closed.");
}
// Create SplitReader.
SplitReader splitReader = splitReaderFactory.get();
int fetcherId = fetcherIdGenerator.getAndIncrement();
SplitFetcher splitFetcher =
new SplitFetcher<>(
fetcherId,
elementsQueue,
splitReader,
errorHandler,
() -> {
fetchers.remove(fetcherId);
// We need this to synchronize status of fetchers to concurrent partners
// as
// ConcurrentHashMap's aggregate status methods including size, isEmpty,
// and
// containsValue are not designed for program control.
elementsQueue.notifyAvailable();
},
this.splitFinishedHook);
fetchers.put(fetcherId, splitFetcher);
return splitFetcher;
}
/**
* Check and shutdown the fetchers that have completed their work.
*
* @return true if all the fetchers have completed the work, false otherwise.
*/
public boolean maybeShutdownFinishedFetchers() {
Iterator>> iter = fetchers.entrySet().iterator();
while (iter.hasNext()) {
Map.Entry> entry = iter.next();
SplitFetcher fetcher = entry.getValue();
if (fetcher.isIdle()) {
LOG.info("Closing splitFetcher {} because it is idle.", entry.getKey());
fetcher.shutdown();
iter.remove();
}
}
return fetchers.isEmpty();
}
/**
* Close the split fetcher manager.
*
* @param timeoutMs the max time in milliseconds to wait.
* @throws Exception when failed to close the split fetcher manager.
*/
public synchronized void close(long timeoutMs) throws Exception {
closed = true;
fetchers.values().forEach(SplitFetcher::shutdown);
executors.shutdown();
if (!executors.awaitTermination(timeoutMs, TimeUnit.MILLISECONDS)) {
LOG.warn(
"Failed to close the source reader in {} ms. There are still {} split fetchers running",
timeoutMs,
fetchers.size());
}
}
public void checkErrors() {
if (uncaughtFetcherException.get() != null) {
throw new RuntimeException(
"One or more fetchers have encountered exception",
uncaughtFetcherException.get());
}
}
// -----------------------
@VisibleForTesting
public int getNumAliveFetchers() {
return fetchers.size();
}
}