SplitAssignment.java
// 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.doris.datasource.split;
import org.apache.doris.common.UserException;
import org.apache.doris.datasource.scan.FederationBackendPolicy;
import org.apache.doris.spi.Split;
import org.apache.doris.system.Backend;
import org.apache.doris.thrift.TScanRangeLocations;
import com.google.common.collect.Multimap;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.io.Closeable;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
/**
* When file splits are supplied in batch mode, splits are generated lazily and assigned in each call of `getNextBatch`.
* `SplitGenerator` provides the file splits, and `FederationBackendPolicy` assigns these splits to backends.
*
* <p>An assignment serves one dispatch of its plan and holds what the backends need from this frontend while they
* scan: the split sources they fetch the splits from, and what the generator opened to produce them (the Flight SQL
* session of a remote Doris scan). {@link #start()} makes the sources reachable to the backends and starts the
* generator, {@link #stop()} releases all of it. The coordinator dispatching the plan starts the assignment and
* stops it when it closes or cancels (ScanNode#start, ScanNode#stopAll), so a plan nobody dispatches holds nothing -
* unless its scan has to start generating splits while it is planned (FileQueryScanNode#needsSampleSplit).
*/
public class SplitAssignment {
private static final Logger LOG = LogManager.getLogger(SplitAssignment.class);
private final FederationBackendPolicy backendPolicy;
private final SplitGenerator splitGenerator;
private final ConcurrentHashMap<Backend, BlockingQueue<Collection<TScanRangeLocations>>> assignment
= new ConcurrentHashMap<>();
private final SplitToScanRange splitToScanRange;
private final Map<String, String> locationProperties;
private final List<String> pathPartitionKeys;
private final boolean fileCacheAdmission;
private final SplitSourceManager splitSourceManager;
private final Object assignLock = new Object();
private Split sampleSplit = null;
private final AtomicBoolean isStopped = new AtomicBoolean(false);
private final AtomicBoolean scheduleFinished = new AtomicBoolean(false);
private UserException exception = null;
// Whether start() ran, the split sources of the backends, and what stop() closes. Guarded by lifecycleLock,
// under which isStopped is set as well: start() may race stop() (a coordinator cancelled while it dispatches the
// plan), and so may addCloseable() (a generator still opening what it hands over when the scan is stopped).
private final Object lifecycleLock = new Object();
private boolean started = false;
private final List<SplitSource> sources = new ArrayList<>();
private final List<Closeable> closeableResources = new ArrayList<>();
public SplitAssignment(
FederationBackendPolicy backendPolicy,
SplitGenerator splitGenerator,
SplitToScanRange splitToScanRange,
Map<String, String> locationProperties,
List<String> pathPartitionKeys,
boolean fileCacheAdmission,
SplitSourceManager splitSourceManager) {
this.backendPolicy = backendPolicy;
this.splitGenerator = splitGenerator;
this.splitToScanRange = splitToScanRange;
this.locationProperties = locationProperties;
this.pathPartitionKeys = pathPartitionKeys;
this.fileCacheAdmission = fileCacheAdmission;
this.splitSourceManager = splitSourceManager;
}
/**
* Makes the split sources reachable to the backends and starts the generator, then waits for its first split
* ({@link #init()}). Called by the coordinator dispatching the plan once the query is admitted (ScanNode#start),
* or while planning by a scan that plans with its first split. Starts once, and a stopped assignment - its
* coordinator was cancelled before it got here - not at all.
*/
public void start() throws UserException {
synchronized (lifecycleLock) {
if (started || isStopped.get()) {
return;
}
started = true;
for (SplitSource source : sources) {
splitSourceManager.registerSplitSource(source);
}
}
init();
}
public void init() throws UserException {
splitGenerator.startSplit(backendPolicy.numBackends());
synchronized (assignLock) {
final int waitIntervalTimeMillis = 100;
final int initTimeoutMillis = 30000; // 30s
int waitTotalTime = 0;
while (sampleSplit == null && needMoreSplit()) {
try {
assignLock.wait(waitIntervalTimeMillis);
} catch (InterruptedException e) {
throw new UserException(e.getMessage(), e);
}
waitTotalTime += waitIntervalTimeMillis;
if (waitTotalTime > initTimeoutMillis) {
throw new UserException("Failed to get first split after waiting for "
+ (waitTotalTime / 1000) + " seconds.");
}
}
}
if (exception != null) {
throw exception;
}
}
public boolean needMoreSplit() {
return !scheduleFinished.get() && !isStopped.get() && exception == null;
}
private void appendBatch(Multimap<Backend, Split> batch) throws UserException {
for (Backend backend : batch.keySet()) {
Collection<Split> splits = batch.get(backend);
List<TScanRangeLocations> locations = new ArrayList<>(splits.size());
for (Split split : splits) {
locations.add(splitToScanRange.getScanRange(backend, locationProperties, split, pathPartitionKeys,
fileCacheAdmission));
}
while (needMoreSplit()) {
BlockingQueue<Collection<TScanRangeLocations>> queue =
assignment.computeIfAbsent(backend, be -> new LinkedBlockingQueue<>(10000));
try {
if (queue.offer(locations, 100, TimeUnit.MILLISECONDS)) {
break;
}
} catch (InterruptedException e) {
addUserException(new UserException("Failed to offer batch split by interrupted", e));
}
}
}
}
/**
* Records the split source a backend fetches its splits from: reachable to the backend from {@link #start()}
* on - at once if the assignment has started already - until {@link #stop()}.
*/
public void registerSource(SplitSource source) {
synchronized (lifecycleLock) {
sources.add(source);
if (started && !isStopped.get()) {
splitSourceManager.registerSplitSource(source);
}
}
}
public Split getSampleSplit() {
return sampleSplit;
}
public void addToQueue(List<Split> splits) throws UserException {
if (splits.isEmpty()) {
return;
}
Multimap<Backend, Split> batch = null;
synchronized (assignLock) {
if (sampleSplit == null) {
sampleSplit = splits.get(0);
assignLock.notify();
}
batch = backendPolicy.computeScanRangeAssignment(splits);
}
appendBatch(batch);
}
private void notifyAssignment() {
synchronized (assignLock) {
assignLock.notify();
}
}
public BlockingQueue<Collection<TScanRangeLocations>> getAssignedSplits(Backend backend) throws UserException {
if (exception != null) {
throw exception;
}
BlockingQueue<Collection<TScanRangeLocations>> splits = assignment.computeIfAbsent(backend,
be -> new LinkedBlockingQueue<>());
if (scheduleFinished.get() && splits.isEmpty() || isStopped.get()) {
return null;
}
return splits;
}
public void setException(UserException e) {
addUserException(e);
notifyAssignment();
}
private void addUserException(UserException e) {
if (exception != null) {
exception.addSuppressed(e);
} else {
exception = e;
}
}
public void finishSchedule() {
scheduleFinished.set(true);
notifyAssignment();
}
/**
* Stops the generator, makes the split sources unreachable to the backends and closes what was handed over
* ({@link #addCloseable}). Idempotent. A failure of the asynchronous split generation is rethrown, once all of
* that is released.
*/
public void stop() {
List<SplitSource> toUnregister;
List<Closeable> toClose;
synchronized (lifecycleLock) {
if (isStopped.get()) {
return;
}
isStopped.set(true);
toUnregister = new ArrayList<>(sources);
toClose = new ArrayList<>(closeableResources);
closeableResources.clear();
}
for (SplitSource source : toUnregister) {
splitSourceManager.removeSplitSource(source.getUniqueId());
}
toClose.forEach(this::closeQuietly);
notifyAssignment();
if (exception != null) {
throw new RuntimeException(exception);
}
}
public boolean isStop() {
return isStopped.get();
}
/**
* Hands over what the backends need until the scan is over, for {@link #stop()} to close: the Flight SQL session
* a remote Doris query runs in. One handed over once the assignment is stopped is closed at once - the generator
* was still opening it when the scan was stopped, and nothing would close it later.
*/
public void addCloseable(Closeable resource) {
synchronized (lifecycleLock) {
if (!isStopped.get()) {
closeableResources.add(resource);
return;
}
}
closeQuietly(resource);
}
private void closeQuietly(Closeable resource) {
try {
resource.close();
} catch (Exception e) {
LOG.warn("close resource error:{}", e.getMessage(), e);
}
}
}