org.jetlinks.rule.engine.defaults.LocalScheduler Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of rule-engine-api Show documentation
There is a newer version: 1.2.2
package org.jetlinks.rule.engine.defaults;

import lombok.Getter;
import lombok.Setter;
import org.jetlinks.rule.engine.api.scheduler.Scheduler;
import org.jetlinks.rule.engine.api.task.Task;
import org.jetlinks.rule.engine.api.worker.Worker;
import org.jetlinks.rule.engine.api.worker.WorkerSelector;
import org.jetlinks.rule.engine.api.scheduler.SchedulingRule;
import org.jetlinks.rule.engine.api.scheduler.ScheduleJob;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.function.Function;

public class LocalScheduler implements Scheduler {

    @Getter
    private final String id;

    @Setter
    private WorkerSelector workerSelector = defaultSelector;

    final static WorkerSelector defaultSelector = (workers1, rule) -> workers1.take(1);

    private final Map workers = new ConcurrentHashMap<>();

    private final Map>> executors = new ConcurrentHashMap<>();

    public LocalScheduler(String id) {
        this.id = id;
    }

    @Override
    public Flux getWorkers() {
        return Flux.fromIterable(workers.values());
    }

    @Override
    public Mono getWorker(String workerId) {

        return Mono.justOrEmpty(workers.get(workerId));
    }

    @Override
    public Mono canSchedule(ScheduleJob job) {
        return findWorker(job.getExecutor(), job)
                .hasElements();
    }

    protected Flux findWorker(String executor, ScheduleJob schedulingRule) {
        return workerSelector
                .select(Flux.fromIterable(workers.values())
                        .filterWhen(exe -> exe.getSupportExecutors()
                                .map(list -> list.contains(executor))
                                .defaultIfEmpty(false)), schedulingRule);
    }

    @Override
    public Flux schedule(ScheduleJob job) {
        //判断调度中的任务
        List tasks = getExecutor(job.getInstanceId(), job.getNodeId());
        if (tasks.isEmpty()) {
            return createExecutor(job);
        }
        return Flux
                .fromIterable(tasks)
                .flatMap(task -> task.setJob(job)
                        .then(task.reload())
                        .thenReturn(task));
    }

    @Override
    public Mono shutdown(String instanceId) {
        return getSchedulingTask(instanceId)
                .flatMap(Task::shutdown)
                .then();
    }

    private Flux createExecutor(ScheduleJob job) {
        return findWorker(job.getExecutor(), job)
                .switchIfEmpty(Mono.error(() -> new UnsupportedOperationException("unsupported executor:" + job.getExecutor())))
                .flatMap(worker -> worker.createTask(id, job))
                .doOnNext(task -> getExecutor(job.getInstanceId(), job.getNodeId()).add(task));
    }

    @Override
    public Flux getSchedulingTask(String instanceId) {
        return Flux.fromIterable(getExecutor(instanceId).values())
                .flatMapIterable(Function.identity());
    }

    @Override
    public Flux getSchedulingTasks() {
        return Flux.fromIterable(executors.values())
                .flatMapIterable(Map::values)
                .flatMapIterable(Function.identity());
    }

    @Override
    public Mono totalTask() {
        return getSchedulingTasks().count();
    }

    private List getExecutor(String instanceId, String nodeId) {
        return getExecutor(instanceId).computeIfAbsent(nodeId, ignore -> new CopyOnWriteArrayList<>());
    }

    private Map> getExecutor(String instanceId) {
        return executors.computeIfAbsent(instanceId, ignore -> new ConcurrentHashMap<>());
    }

    public void addWorker(Worker worker) {
        this.workers.put(worker.getId(), worker);
    }
}