wsjz commented on code in PR #26356: URL: https://github.com/apache/doris/pull/26356#discussion_r1427703299
########## fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/load/LoadMgr.java: ########## @@ -0,0 +1,340 @@ +// 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.nereids.jobs.load; + +import org.apache.doris.analysis.CompoundPredicate; +import org.apache.doris.catalog.Database; +import org.apache.doris.catalog.Env; +import org.apache.doris.common.AnalysisException; +import org.apache.doris.common.CaseSensibility; +import org.apache.doris.common.DdlException; +import org.apache.doris.common.LabelAlreadyUsedException; +import org.apache.doris.common.PatternMatcher; +import org.apache.doris.common.PatternMatcherWrapper; +import org.apache.doris.common.util.LogBuilder; +import org.apache.doris.common.util.LogKey; +import org.apache.doris.job.base.JobExecuteType; +import org.apache.doris.job.base.JobExecutionConfiguration; +import org.apache.doris.job.common.JobStatus; +import org.apache.doris.job.exception.JobException; +import org.apache.doris.job.extensions.insert.InsertJob; +import org.apache.doris.job.manager.JobManager; +import org.apache.doris.load.loadv2.JobState; +import org.apache.doris.nereids.jobs.load.replay.ReplayLoadLog; + +import com.google.common.base.Strings; +import com.google.common.collect.Lists; +import org.apache.commons.lang3.StringUtils; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.stream.Collectors; + +/** + * load manager + */ +public class LoadMgr { + private static final Logger LOG = LogManager.getLogger(LoadMgr.class); + private Map<Long, InsertJob> loadIdToJob = new HashMap<>(); + private final Map<Long, Map<String, List<InsertJob>>> dbIdToLabelToLoadJobs = new ConcurrentHashMap<>(); + + // lock for export job + // lock is private and must use after db lock + private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(true); + + private void readLock() { + lock.readLock().lock(); + } + + private void readUnlock() { + lock.readLock().unlock(); + } + + private void writeLock() { + lock.writeLock().lock(); + } + + private void writeUnlock() { + lock.writeLock().unlock(); + } + + private JobManager<InsertJob, ?> getJobManager() { + return Env.getCurrentEnv().getJobManager(); + } + + /** + * add load job and add tasks + * @param loadJob job + */ + public void addLoadJob(InsertJob loadJob) throws DdlException { + writeLock(); + try { + Map<String, List<InsertJob>> labelToLoadJobs = dbIdToLabelToLoadJobs.get(loadJob.getDbId()); + if (labelToLoadJobs != null && labelToLoadJobs.containsKey(loadJob.getLabel())) { + throw new LabelAlreadyUsedException(loadJob.getLabel()); + } + unprotectAddJob(loadJob); + } catch (LabelAlreadyUsedException e) { + throw new RuntimeException(e); + } finally { + writeUnlock(); + } + } + + private void unprotectAddJob(InsertJob job) throws DdlException { + loadIdToJob.put(job.getJobId(), job); + try { + getJobManager().registerJob(job); + if (!dbIdToLabelToLoadJobs.containsKey(job.getDbId())) { + dbIdToLabelToLoadJobs.put(job.getDbId(), new ConcurrentHashMap<>()); + } + Map<String, List<InsertJob>> labelToLoadJobs = dbIdToLabelToLoadJobs.get(job.getDbId()); + if (!labelToLoadJobs.containsKey(job.getLabel())) { + labelToLoadJobs.put(job.getLabel(), new ArrayList<>()); + } + labelToLoadJobs.get(job.getLabel()).add(job); + } catch (org.apache.doris.job.exception.JobException e) { + throw new DdlException(e.getMessage(), e); + } + } + + /** + * replay load job + * @param replayLoadLog load log + * @throws DdlException ex + */ + public void replayLoadJob(ReplayLoadLog replayLoadLog) throws DdlException { + writeLock(); + try { + if (replayLoadLog instanceof ReplayLoadLog.ReplayCreateLoadLog) { + InsertJob loadJob = new InsertJob((ReplayLoadLog.ReplayCreateLoadLog) replayLoadLog); + JobExecutionConfiguration jobConfig = new JobExecutionConfiguration(); + jobConfig.setExecuteType(JobExecuteType.INSTANT); + loadJob.setJobConfig(jobConfig); + addLoadJob(loadJob); + LOG.info(new LogBuilder(LogKey.LOAD_JOB, loadJob.getJobId()).add("msg", "replay create load job") + .build()); + } else if (replayLoadLog instanceof ReplayLoadLog.ReplayEndLoadLog) { + InsertJob job = loadIdToJob.get(replayLoadLog.getId()); + if (job == null) { + // This should not happen. + // Last time I found that when user submit a job with already used label, an END_LOAD_JOB edit log + // will be written but the job is not added to 'idToLoadJob', so this job here we got will be null. + // And this bug has been fixed. + // Just add a log here to observe. + LOG.warn("job does not exist when replaying end load job edit log: {}", replayLoadLog); + return; + } + job.unprotectReadEndOperation((ReplayLoadLog.ReplayEndLoadLog) replayLoadLog); + LOG.info(new LogBuilder(LogKey.LOAD_JOB, replayLoadLog.getId()).add("operation", replayLoadLog) + .add("msg", "replay end load job").build()); + } else { + throw new DdlException("Unsupported replay job type. "); + } + } finally { + writeUnlock(); + } + } + + // public void updateJobProgress(Long jobId, Long beId, TUniqueId loadId, TUniqueId fragmentId, long scannedRows, + // long scannedBytes, boolean isDone) { + // LoadJobExecutor job = loadIdToJob.get(jobId); + // if (job != null) { + // job.updateLoadingStatus(beId, loadId, fragmentId, scannedRows, scannedBytes, isDone); + // } + // } + + /** + * cancel job + * + * @param dbName dbName + * @param label job label + * @param state job state + * @param operator filter operator, like or equals + */ + public void cancelLoadJob(String dbName, String label, String state, CompoundPredicate.Operator operator) + throws JobException, AnalysisException, DdlException { + Database db = Env.getCurrentInternalCatalog().getDbOrDdlException(dbName); + // List of load jobs waiting to be cancelled + List<InsertJob> uncompletedLoadJob; + readLock(); + try { + Map<String, List<InsertJob>> labelToLoadJobs = dbIdToLabelToLoadJobs.get(db.getId()); + if (labelToLoadJobs == null) { + throw new JobException("Load job does not exist"); + } + List<InsertJob> matchLoadJobs = Lists.newArrayList(); + addNeedCancelLoadJob(label, state, operator, + labelToLoadJobs.values().stream().flatMap(Collection::stream).collect(Collectors.toList()), + matchLoadJobs); + if (matchLoadJobs.isEmpty()) { + throw new JobException("Load job does not exist"); + } + // check state here + uncompletedLoadJob = + matchLoadJobs.stream().filter(InsertJob::isRunning) + .collect(Collectors.toList()); + if (uncompletedLoadJob.isEmpty()) { + throw new JobException("There is no uncompleted job"); + } + } finally { + readUnlock(); + } + for (InsertJob loadJob : uncompletedLoadJob) { + try { + loadJob.cancelJob(); + } catch (JobException e) { + LOG.warn("Fail to cancel job, its label: {}", loadJob.getLabel()); + } + } + } + + private static void addNeedCancelLoadJob(String label, String state, Review Comment: not required and I suggest make a method static if it's not use any instance fields. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org For additional commands, e-mail: commits-h...@doris.apache.org
