| /* |
| * 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.nemo.runtime.master.scheduler; |
| |
| import com.google.common.annotations.VisibleForTesting; |
| import org.apache.nemo.common.Pair; |
| import org.apache.nemo.runtime.common.plan.Task; |
| import org.apache.nemo.runtime.master.resource.ExecutorRepresenter; |
| import org.apache.reef.annotations.audience.DriverSide; |
| |
| import javax.annotation.concurrent.ThreadSafe; |
| import javax.inject.Inject; |
| import java.util.HashMap; |
| import java.util.Map; |
| import java.util.Optional; |
| import java.util.Set; |
| import java.util.function.BiFunction; |
| import java.util.function.Consumer; |
| import java.util.stream.Collectors; |
| |
| /** |
| * (WARNING) This class must be thread-safe. |
| * Maintains map between executor id and {@link ExecutorRepresenter}. |
| */ |
| @DriverSide |
| @ThreadSafe |
| public final class ExecutorRegistry { |
| /** |
| * States of an executor. |
| */ |
| enum ExecutorState { |
| RUNNING, |
| FAILED, |
| TERMINATED |
| } |
| |
| private final Map<String, Pair<ExecutorRepresenter, ExecutorState>> executors; |
| |
| @Inject |
| private ExecutorRegistry() { |
| this.executors = new HashMap<>(); |
| } |
| |
| /** |
| * Static constructor for manual usage. |
| * @return a new instance of ExecutorRegistry. |
| */ |
| public static ExecutorRegistry newInstance() { |
| return new ExecutorRegistry(); |
| } |
| |
| synchronized void registerExecutor(final ExecutorRepresenter executor) { |
| final String executorId = executor.getExecutorId(); |
| if (executors.containsKey(executorId)) { |
| throw new IllegalArgumentException("Duplicate executor: " + executor.toString()); |
| } else { |
| executors.put(executorId, Pair.of(executor, ExecutorState.RUNNING)); |
| } |
| } |
| |
| public synchronized void viewExecutors(final Consumer<Set<ExecutorRepresenter>> consumer) { |
| consumer.accept(getRunningExecutors()); |
| } |
| |
| public synchronized int getNumberOfRunningExecutors() { |
| return getRunningExecutors().size(); |
| } |
| |
| synchronized void updateExecutor( |
| final String executorId, |
| final BiFunction<ExecutorRepresenter, ExecutorState, Pair<ExecutorRepresenter, ExecutorState>> updater) { |
| final Pair<ExecutorRepresenter, ExecutorState> pair = executors.get(executorId); |
| if (pair == null) { |
| throw new IllegalArgumentException("Unknown executor id " + executorId); |
| } else { |
| executors.put(executorId, updater.apply(pair.left(), pair.right())); |
| } |
| } |
| |
| synchronized void terminate() { |
| for (final ExecutorRepresenter executor : getRunningExecutors()) { |
| executor.shutDown(); |
| executors.put(executor.getExecutorId(), Pair.of(executor, ExecutorState.TERMINATED)); |
| } |
| } |
| |
| /** |
| * Retrieves the executor to which the given task was scheduled. |
| * |
| * @param taskId of the task to search. |
| * @return the {@link ExecutorRepresenter} of the executor the task was scheduled to. |
| */ |
| @VisibleForTesting |
| synchronized Optional<ExecutorRepresenter> findExecutorForTask(final String taskId) { |
| for (final ExecutorRepresenter executor : getRunningExecutors()) { |
| for (final Task runningTask : executor.getRunningTasks()) { |
| if (runningTask.getTaskId().equals(taskId)) { |
| return Optional.of(executor); |
| } |
| } |
| } |
| return Optional.empty(); |
| } |
| |
| private Set<ExecutorRepresenter> getRunningExecutors() { |
| return executors.values() |
| .stream() |
| .filter(pair -> pair.right().equals(ExecutorState.RUNNING)) |
| .map(Pair::left) |
| .collect(Collectors.toSet()); |
| } |
| |
| @Override |
| public String toString() { |
| return executors.toString(); |
| } |
| } |