org.apache.flink.runtime.source.coordinator.SourceCoordinatorContext 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.runtime.source.coordinator;
import org.apache.flink.annotation.Internal;
import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.api.connector.source.ReaderInfo;
import org.apache.flink.api.connector.source.SourceEvent;
import org.apache.flink.api.connector.source.SourceSplit;
import org.apache.flink.api.connector.source.SplitEnumerator;
import org.apache.flink.api.connector.source.SplitEnumeratorContext;
import org.apache.flink.api.connector.source.SplitsAssignment;
import org.apache.flink.api.connector.source.SupportsIntermediateNoMoreSplits;
import org.apache.flink.core.io.SimpleVersionedSerializer;
import org.apache.flink.metrics.groups.SplitEnumeratorMetricGroup;
import org.apache.flink.runtime.operators.coordination.OperatorCoordinator;
import org.apache.flink.runtime.operators.coordination.OperatorEvent;
import org.apache.flink.runtime.source.event.AddSplitEvent;
import org.apache.flink.runtime.source.event.NoMoreSplitsEvent;
import org.apache.flink.runtime.source.event.SourceEventWrapper;
import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.FlinkRuntimeException;
import org.apache.flink.util.ThrowableCatchingRunnable;
import org.apache.flink.util.concurrent.ExecutorThreadFactory;
import org.apache.flink.shaded.guava30.com.google.common.collect.Iterables;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;
import static org.apache.flink.runtime.operators.coordination.ComponentClosingUtils.shutdownExecutorForcefully;
import static org.apache.flink.util.Preconditions.checkState;
/**
* A context class for the {@link OperatorCoordinator}. Compared with {@link SplitEnumeratorContext}
* this class allows interaction with state and sending {@link OperatorEvent} to the SourceOperator
* while {@link SplitEnumeratorContext} only allows sending {@link SourceEvent}.
*
* The context serves a few purposes:
*
*
* - Information provider - The context provides necessary information to the enumerator for it
* to know what is the status of the source readers and their split assignments. These
* information allows the split enumerator to do the coordination.
*
- Action taker - The context also provides a few actions that the enumerator can take to
* carry out the coordination. So far there are two actions: 1) assign splits to the source
* readers. and 2) sens a custom {@link SourceEvent SourceEvents} to the source readers.
*
- Thread model enforcement - The context ensures that all the manipulations to the
* coordinator state are handled by the same thread.
*
*
* @param the type of the splits.
*/
@Internal
public class SourceCoordinatorContext
implements SplitEnumeratorContext, SupportsIntermediateNoMoreSplits, AutoCloseable {
private static final Logger LOG = LoggerFactory.getLogger(SourceCoordinatorContext.class);
private final ScheduledExecutorService workerExecutor;
private final ScheduledExecutorService coordinatorExecutor;
private final ExecutorNotifier notifier;
private final OperatorCoordinator.Context operatorCoordinatorContext;
private final SimpleVersionedSerializer splitSerializer;
private final ConcurrentMap> registeredReaders;
private final SplitAssignmentTracker assignmentTracker;
private final SourceCoordinatorProvider.CoordinatorExecutorThreadFactory
coordinatorThreadFactory;
private final SubtaskGateways subtaskGateways;
private final String coordinatorThreadName;
private final boolean supportsConcurrentExecutionAttempts;
private final boolean[] subtaskHasNoMoreSplits;
private volatile boolean closed;
public SourceCoordinatorContext(
SourceCoordinatorProvider.CoordinatorExecutorThreadFactory coordinatorThreadFactory,
int numWorkerThreads,
OperatorCoordinator.Context operatorCoordinatorContext,
SimpleVersionedSerializer splitSerializer,
boolean supportsConcurrentExecutionAttempts) {
this(
Executors.newScheduledThreadPool(1, coordinatorThreadFactory),
Executors.newScheduledThreadPool(
numWorkerThreads,
new ExecutorThreadFactory(
coordinatorThreadFactory.getCoordinatorThreadName() + "-worker")),
coordinatorThreadFactory,
operatorCoordinatorContext,
splitSerializer,
new SplitAssignmentTracker<>(),
supportsConcurrentExecutionAttempts);
}
// Package private method for unit test.
@VisibleForTesting
SourceCoordinatorContext(
ScheduledExecutorService coordinatorExecutor,
ScheduledExecutorService workerExecutor,
SourceCoordinatorProvider.CoordinatorExecutorThreadFactory coordinatorThreadFactory,
OperatorCoordinator.Context operatorCoordinatorContext,
SimpleVersionedSerializer splitSerializer,
SplitAssignmentTracker splitAssignmentTracker,
boolean supportsConcurrentExecutionAttempts) {
this.workerExecutor = workerExecutor;
this.coordinatorExecutor = coordinatorExecutor;
this.coordinatorThreadFactory = coordinatorThreadFactory;
this.operatorCoordinatorContext = operatorCoordinatorContext;
this.splitSerializer = splitSerializer;
this.registeredReaders = new ConcurrentHashMap<>();
this.assignmentTracker = splitAssignmentTracker;
this.coordinatorThreadName = coordinatorThreadFactory.getCoordinatorThreadName();
this.supportsConcurrentExecutionAttempts = supportsConcurrentExecutionAttempts;
final int parallelism = operatorCoordinatorContext.currentParallelism();
this.subtaskGateways = new SubtaskGateways(parallelism);
this.subtaskHasNoMoreSplits = new boolean[parallelism];
Arrays.fill(subtaskHasNoMoreSplits, false);
final Executor errorHandlingCoordinatorExecutor =
(runnable) ->
coordinatorExecutor.execute(
new ThrowableCatchingRunnable(
this::handleUncaughtExceptionFromAsyncCall, runnable));
this.notifier = new ExecutorNotifier(workerExecutor, errorHandlingCoordinatorExecutor);
}
boolean isConcurrentExecutionAttemptsSupported() {
return supportsConcurrentExecutionAttempts;
}
@Override
public SplitEnumeratorMetricGroup metricGroup() {
return null;
}
@Override
public void sendEventToSourceReader(int subtaskId, SourceEvent event) {
checkState(
!supportsConcurrentExecutionAttempts,
"The split enumerator must invoke SplitEnumeratorContext"
+ "#sendEventToSourceReader(int, int, SourceEvent) instead of "
+ "SplitEnumeratorContext#sendEventToSourceReader(int, SourceEvent) "
+ "to send custom source events in concurrent execution attempts scenario "
+ "(e.g. if speculative execution is enabled).");
checkSubtaskIndex(subtaskId);
callInCoordinatorThread(
() -> {
final OperatorCoordinator.SubtaskGateway gateway =
subtaskGateways.getOnlyGatewayAndCheckReady(subtaskId);
gateway.sendEvent(new SourceEventWrapper(event));
return null;
},
String.format("Failed to send event %s to subtask %d", event, subtaskId));
}
@Override
public void sendEventToSourceReader(int subtaskId, int attemptNumber, SourceEvent event) {
checkSubtaskIndex(subtaskId);
callInCoordinatorThread(
() -> {
final OperatorCoordinator.SubtaskGateway gateway =
subtaskGateways.getGatewayAndCheckReady(subtaskId, attemptNumber);
gateway.sendEvent(new SourceEventWrapper(event));
return null;
},
String.format(
"Failed to send event %s to subtask %d (#%d)",
event, subtaskId, attemptNumber));
}
void sendEventToSourceOperator(int subtaskId, OperatorEvent event) {
checkSubtaskIndex(subtaskId);
callInCoordinatorThread(
() -> {
final OperatorCoordinator.SubtaskGateway gateway =
subtaskGateways.getOnlyGatewayAndCheckReady(subtaskId);
gateway.sendEvent(event);
return null;
},
String.format("Failed to send event %s to subtask %d", event, subtaskId));
}
void sendEventToSourceOperatorIfTaskReady(int subtaskId, OperatorEvent event) {
checkSubtaskIndex(subtaskId);
callInCoordinatorThread(
() -> {
final OperatorCoordinator.SubtaskGateway gateway =
subtaskGateways.getOnlyGatewayAndNotCheckReady(subtaskId);
if (gateway != null) {
gateway.sendEvent(event);
}
return null;
},
String.format("Failed to send event %s to subtask %d", event, subtaskId));
}
@Override
public int currentParallelism() {
return operatorCoordinatorContext.currentParallelism();
}
@Override
public Map registeredReaders() {
final Map readers = new HashMap<>();
for (Map.Entry> entry :
registeredReaders.entrySet()) {
final int subtaskIndex = entry.getKey();
final Map attemptReaders = entry.getValue();
int earliestAttempt = Integer.MAX_VALUE;
for (int attemptNumber : attemptReaders.keySet()) {
if (attemptNumber < earliestAttempt) {
earliestAttempt = attemptNumber;
}
}
readers.put(subtaskIndex, attemptReaders.get(earliestAttempt));
}
return Collections.unmodifiableMap(readers);
}
@Override
public Map> registeredReadersOfAttempts() {
return Collections.unmodifiableMap(registeredReaders);
}
@Override
public void assignSplits(SplitsAssignment assignment) {
// Ensure the split assignment is done by the coordinator executor.
callInCoordinatorThread(
() -> {
// Ensure all the subtasks in the assignment have registered.
assignment
.assignment()
.forEach(
(id, splits) -> {
if (!registeredReaders.containsKey(id)) {
throw new IllegalArgumentException(
String.format(
"Cannot assign splits %s to subtask %d because the subtask is not registered.",
splits, id));
}
});
assignmentTracker.recordSplitAssignment(assignment);
assignSplitsToAttempts(assignment);
return null;
},
String.format("Failed to assign splits %s due to ", assignment));
}
@Override
public void signalNoMoreSplits(int subtask) {
checkSubtaskIndex(subtask);
// Ensure the split assignment is done by the coordinator executor.
callInCoordinatorThread(
() -> {
subtaskHasNoMoreSplits[subtask] = true;
signalNoMoreSplitsToAttempts(subtask);
return null; // void return value
},
"Failed to send 'NoMoreSplits' to reader " + subtask);
}
@Override
public void signalIntermediateNoMoreSplits(int subtask) {
checkSubtaskIndex(subtask);
// It's an intermediate noMoreSplit event, notify subtask to deal with this event.
callInCoordinatorThread(
() -> {
signalNoMoreSplitsToAttempts(subtask);
return null;
},
"Failed to send 'IntermediateNoMoreSplits' to reader " + subtask);
}
@Override
public void callAsync(
Callable callable,
BiConsumer handler,
long initialDelay,
long period) {
notifier.notifyReadyAsync(callable, handler, initialDelay, period);
}
@Override
public void callAsync(Callable callable, BiConsumer handler) {
notifier.notifyReadyAsync(callable, handler);
}
/** {@inheritDoc} If the runnable throws an Exception, the corresponding job is failed. */
@Override
public void runInCoordinatorThread(Runnable runnable) {
// when using a ScheduledThreadPool, uncaught exception handler catches only
// exceptions thrown by the threadPool, so manually call it when the exception is
// thrown by the runnable
coordinatorExecutor.execute(
new ThrowableCatchingRunnable(
throwable ->
coordinatorThreadFactory.uncaughtException(
Thread.currentThread(), throwable),
runnable));
}
@Override
public void close() throws InterruptedException {
closed = true;
// Close quietly so the closing sequence will be executed completely.
shutdownExecutorForcefully(workerExecutor, Duration.ofNanos(Long.MAX_VALUE));
shutdownExecutorForcefully(coordinatorExecutor, Duration.ofNanos(Long.MAX_VALUE));
}
@VisibleForTesting
boolean isClosed() {
return closed;
}
// --------- Package private additional methods for the SourceCoordinator ------------
void attemptReady(OperatorCoordinator.SubtaskGateway gateway) {
checkState(coordinatorThreadFactory.isCurrentThreadCoordinatorThread());
subtaskGateways.registerSubtaskGateway(gateway);
}
void attemptFailed(int subtaskIndex, int attemptNumber) {
checkState(coordinatorThreadFactory.isCurrentThreadCoordinatorThread());
subtaskGateways.unregisterSubtaskGateway(subtaskIndex, attemptNumber);
}
void subtaskReset(int subtaskIndex) {
checkState(coordinatorThreadFactory.isCurrentThreadCoordinatorThread());
subtaskGateways.reset(subtaskIndex);
registeredReaders.remove(subtaskIndex);
subtaskHasNoMoreSplits[subtaskIndex] = false;
}
boolean hasNoMoreSplits(int subtaskIndex) {
return subtaskHasNoMoreSplits[subtaskIndex];
}
/**
* Fail the job with the given cause.
*
* @param cause the cause of the job failure.
*/
void failJob(Throwable cause) {
operatorCoordinatorContext.failJob(cause);
}
void handleUncaughtExceptionFromAsyncCall(Throwable t) {
if (closed) {
return;
}
ExceptionUtils.rethrowIfFatalErrorOrOOM(t);
LOG.error(
"Exception while handling result from async call in {}. Triggering job failover.",
coordinatorThreadName,
t);
failJob(t);
}
/**
* Behavior of SourceCoordinatorContext on checkpoint.
*
* @param checkpointId The id of the ongoing checkpoint.
*/
void onCheckpoint(long checkpointId) throws Exception {
assignmentTracker.onCheckpoint(checkpointId);
}
/**
* Register a source reader.
*
* @param subtaskId the subtask id of the source reader.
* @param attemptNumber the attempt number of the source reader.
* @param location the location of the source reader.
*/
void registerSourceReader(int subtaskId, int attemptNumber, String location) {
final Map attemptReaders =
registeredReaders.computeIfAbsent(subtaskId, k -> new ConcurrentHashMap<>());
checkState(
!attemptReaders.containsKey(attemptNumber),
"ReaderInfo of subtask %s (#%s) already exists.",
subtaskId,
attemptNumber);
attemptReaders.put(attemptNumber, new ReaderInfo(subtaskId, location));
sendCachedSplitsToNewlyRegisteredReader(subtaskId, attemptNumber);
}
/**
* Unregister a source reader.
*
* @param subtaskId the subtask id of the source reader.
* @param attemptNumber the attempt number of the source reader.
*/
void unregisterSourceReader(int subtaskId, int attemptNumber) {
final Map attemptReaders = registeredReaders.get(subtaskId);
if (attemptReaders != null) {
attemptReaders.remove(attemptNumber);
if (attemptReaders.isEmpty()) {
registeredReaders.remove(subtaskId);
}
}
}
/**
* Get the split to put back. This only happens when a source reader subtask has failed.
*
* @param subtaskId the failed subtask id.
* @param restoredCheckpointId the checkpoint that the task is recovered to.
* @return A list of splits that needs to be added back to the {@link SplitEnumerator}.
*/
List getAndRemoveUncheckpointedAssignment(int subtaskId, long restoredCheckpointId) {
return assignmentTracker.getAndRemoveUncheckpointedAssignment(
subtaskId, restoredCheckpointId);
}
/**
* Invoked when a successful checkpoint has been taken.
*
* @param checkpointId the id of the successful checkpoint.
*/
void onCheckpointComplete(long checkpointId) {
assignmentTracker.onCheckpointComplete(checkpointId);
}
OperatorCoordinator.Context getCoordinatorContext() {
return operatorCoordinatorContext;
}
// ---------------- Executor methods to avoid use coordinatorExecutor directly -----------------
Future submitTask(Runnable task) {
return coordinatorExecutor.submit(task);
}
/** To avoid period task lost, we should handle the potential exception throw by task. */
ScheduledFuture schedulePeriodTask(
Runnable command, long initDelay, long period, TimeUnit unit) {
return coordinatorExecutor.scheduleAtFixedRate(
() -> {
try {
command.run();
} catch (Throwable t) {
handleUncaughtExceptionFromAsyncCall(t);
}
},
initDelay,
period,
unit);
}
// ---------------- private helper methods -----------------
private void checkSubtaskIndex(int subtaskIndex) {
if (subtaskIndex < 0 || subtaskIndex >= getCoordinatorContext().currentParallelism()) {
throw new IllegalArgumentException(
String.format(
"Subtask index %d is out of bounds [0, %s)",
subtaskIndex, getCoordinatorContext().currentParallelism()));
}
}
/**
* A helper method that delegates the callable to the coordinator thread if the current thread
* is not the coordinator thread, otherwise call the callable right away.
*
* @param callable the callable to delegate.
*/
private V callInCoordinatorThread(Callable callable, String errorMessage) {
// Ensure the split assignment is done by the coordinator executor.
if (!coordinatorThreadFactory.isCurrentThreadCoordinatorThread()) {
try {
final Callable guardedCallable =
() -> {
try {
return callable.call();
} catch (Throwable t) {
LOG.error("Uncaught Exception in Source Coordinator Executor", t);
ExceptionUtils.rethrowException(t);
return null;
}
};
return coordinatorExecutor.submit(guardedCallable).get();
} catch (InterruptedException | ExecutionException e) {
throw new FlinkRuntimeException(errorMessage, e);
}
}
try {
return callable.call();
} catch (Throwable t) {
LOG.error("Uncaught Exception in Source Coordinator Executor", t);
throw new FlinkRuntimeException(errorMessage, t);
}
}
private void assignSplitsToAttempts(SplitsAssignment assignment) {
assignment.assignment().forEach((index, splits) -> assignSplitsToAttempts(index, splits));
}
private void assignSplitsToAttempts(int subtaskIndex, List splits) {
getRegisteredAttempts(subtaskIndex)
.forEach(attempt -> assignSplitsToAttempt(subtaskIndex, attempt, splits));
}
private void assignSplitsToAttempt(int subtaskIndex, int attemptNumber, List splits) {
if (splits.isEmpty()) {
return;
}
checkAttemptReaderReady(subtaskIndex, attemptNumber);
final AddSplitEvent addSplitEvent;
try {
addSplitEvent = new AddSplitEvent<>(splits, splitSerializer);
} catch (IOException e) {
throw new FlinkRuntimeException("Failed to serialize splits.", e);
}
final OperatorCoordinator.SubtaskGateway gateway =
subtaskGateways.getGatewayAndCheckReady(subtaskIndex, attemptNumber);
gateway.sendEvent(addSplitEvent);
}
private void signalNoMoreSplitsToAttempts(int subtaskIndex) {
getRegisteredAttempts(subtaskIndex)
.forEach(attemptNumber -> signalNoMoreSplitsToAttempt(subtaskIndex, attemptNumber));
}
private void signalNoMoreSplitsToAttempt(int subtaskIndex, int attemptNumber) {
checkAttemptReaderReady(subtaskIndex, attemptNumber);
final OperatorCoordinator.SubtaskGateway gateway =
subtaskGateways.getGatewayAndCheckReady(subtaskIndex, attemptNumber);
gateway.sendEvent(new NoMoreSplitsEvent());
}
private void checkAttemptReaderReady(int subtaskIndex, int attemptNumber) {
checkState(registeredReaders.containsKey(subtaskIndex));
checkState(getRegisteredAttempts(subtaskIndex).contains(attemptNumber));
}
private Set getRegisteredAttempts(int subtaskIndex) {
return registeredReaders.get(subtaskIndex).keySet();
}
private void sendCachedSplitsToNewlyRegisteredReader(int subtaskIndex, int attemptNumber) {
// For batch jobs, checkpoints will never happen so that the un-checkpointed assignments in
// assignmentTracker can be seen as cached splits. For streaming jobs,
// #supportsConcurrentExecutionAttempts should be false and un-checkpointed assignments
// should be empty when a new reader is registered (cleared when the last reader failed).
final LinkedHashSet cachedSplits =
assignmentTracker.uncheckpointedAssignments().get(subtaskIndex);
if (cachedSplits != null) {
if (supportsConcurrentExecutionAttempts) {
assignSplitsToAttempt(subtaskIndex, attemptNumber, new ArrayList<>(cachedSplits));
if (hasNoMoreSplits(subtaskIndex)) {
signalNoMoreSplitsToAttempt(subtaskIndex, attemptNumber);
}
} else {
throw new IllegalStateException("No cached split is expected.");
}
}
}
/** Maintains the subtask gateways for different execution attempts of different subtasks. */
private static class SubtaskGateways {
private final Map[] gateways;
private SubtaskGateways(int parallelism) {
gateways = new Map[parallelism];
for (int i = 0; i < parallelism; i++) {
gateways[i] = new HashMap<>();
}
}
private void registerSubtaskGateway(OperatorCoordinator.SubtaskGateway gateway) {
final int subtaskIndex = gateway.getSubtask();
final int attemptNumber = gateway.getExecution().getAttemptNumber();
checkState(
!gateways[subtaskIndex].containsKey(attemptNumber),
"Already have a subtask gateway for %s (#%s).",
subtaskIndex,
attemptNumber);
gateways[subtaskIndex].put(attemptNumber, gateway);
}
private void unregisterSubtaskGateway(int subtaskIndex, int attemptNumber) {
gateways[subtaskIndex].remove(attemptNumber);
}
private OperatorCoordinator.SubtaskGateway getOnlyGatewayAndCheckReady(int subtaskIndex) {
checkState(
gateways[subtaskIndex].size() > 0,
"Subtask %s is not ready yet to receive events.",
subtaskIndex);
return Iterables.getOnlyElement(gateways[subtaskIndex].values());
}
private OperatorCoordinator.SubtaskGateway getOnlyGatewayAndNotCheckReady(
int subtaskIndex) {
if (gateways[subtaskIndex].size() > 0) {
return Iterables.getOnlyElement(gateways[subtaskIndex].values());
} else {
return null;
}
}
private OperatorCoordinator.SubtaskGateway getGatewayAndCheckReady(
int subtaskIndex, int attemptNumber) {
final OperatorCoordinator.SubtaskGateway gateway =
gateways[subtaskIndex].get(attemptNumber);
if (gateway != null) {
return gateway;
}
throw new IllegalStateException(
String.format(
"Subtask %d (#%d) is not ready yet to receive events.",
subtaskIndex, attemptNumber));
}
private void reset(int subtaskIndex) {
gateways[subtaskIndex].clear();
}
}
}