Source code

001/**
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *   http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing,
013 * software distributed under the License is distributed on an
014 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
015 * KIND, either express or implied.  See the License for the
016 * specific language governing permissions and limitations
017 * under the License.
018 */
019package org.apache.reef.runtime.common.driver;
020
021import org.apache.reef.driver.parameters.DriverRestartHandler;
022import org.apache.reef.tang.annotations.Parameter;
023import org.apache.reef.util.Optional;
024import org.apache.reef.wake.EventHandler;
025import org.apache.reef.wake.time.event.StartTime;
026
027import javax.inject.Inject;
028import java.util.Set;
029import java.util.logging.Level;
030import java.util.logging.Logger;
031
032/**
033 * This is bound to the start event of the clock and dispatches it to the approriate application code.
034 */
035public final class DriverStartHandler implements EventHandler<StartTime> {
036  private static final Logger LOG = Logger.getLogger(DriverStartHandler.class.getName());
037
038  private final Set<EventHandler<StartTime>> startHandlers;
039  private final Optional<EventHandler<StartTime>> restartHandler;
040  private final DriverStatusManager driverStatusManager;
041
042  @Inject
043  DriverStartHandler(final @Parameter(org.apache.reef.driver.parameters.DriverStartHandler.class) Set<EventHandler<StartTime>> startHandler,
044                     final @Parameter(DriverRestartHandler.class) EventHandler<StartTime> restartHandler,
045                     final DriverStatusManager driverStatusManager) {
046    this.startHandlers = startHandler;
047    this.restartHandler = Optional.of(restartHandler);
048    this.driverStatusManager = driverStatusManager;
049    LOG.log(Level.FINE, "Instantiated `DriverStartHandler with StartHandler [{0}] and RestartHandler [{1}]",
050        new String[]{this.startHandlers.toString(), this.restartHandler.toString()});
051  }
052
053  @Inject
054  DriverStartHandler(final @Parameter(org.apache.reef.driver.parameters.DriverStartHandler.class) Set<EventHandler<StartTime>> startHandler,
055                     final DriverStatusManager driverStatusManager) {
056    this.startHandlers = startHandler;
057    this.restartHandler = Optional.empty();
058    this.driverStatusManager = driverStatusManager;
059    LOG.log(Level.FINE, "Instantiated `DriverStartHandler with StartHandler [{0}] and no RestartHandler",
060        this.startHandlers.toString());
061  }
062
063  @Override
064  public void onNext(final StartTime startTime) {
065    if (isRestart()) {
066      this.onRestart(startTime);
067    } else {
068      this.onStart(startTime);
069    }
070  }
071
072  private void onRestart(final StartTime startTime) {
073    if (restartHandler.isPresent()) {
074      this.restartHandler.get().onNext(startTime);
075    } else {
076      // TODO: We might have to indicate this to YARN somehow such that it doesn't try another time.
077      throw new RuntimeException("Driver restart happened, but no ON_DRIVER_RESTART handler is bound.");
078    }
079  }
080
081  private void onStart(final StartTime startTime) {
082    for (final EventHandler<StartTime> startHandler : this.startHandlers) {
083      startHandler.onNext(startTime);
084    }
085  }
086
087  /**
088   * @return true, if the Driver is in fact being restarted.
089   */
090  private boolean isRestart() {
091    return this.driverStatusManager.getNumPreviousContainers() > 0;
092  }
093}