* 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.

*/

* ZooKeeper-based implementation of {@link CentralizedServiceWorker}.

*

* @param <I> Vertex id

* @param <V> Vertex data

* @param <E> Edge data

*/

    V extends Writable, E extends Writable>

    extends BspService<I, V, E>

    implements CentralizedServiceWorker<I, V, E>,

    ResetSuperstepMetricsObserver {

  /** Name of gauge for time spent waiting on other workers */

  public static final String TIMER_WAIT_REQUESTS = "wait-requests-us";

  /** Class logger */

  private static final Logger LOG = Logger.getLogger(BspServiceWorker.class);

  /** My process health znode */

  private String myHealthZnode;

  /** Worker info */

  private final WorkerInfo workerInfo;

  /** Worker graph partitioner */

  private final WorkerGraphPartitioner<I, V, E> workerGraphPartitioner;

  /** Local Data for each worker */

  private final LocalData<I, V, E, ? extends Writable> localData;

  /** Used to translate Edges during vertex input phase based on localData */

  private final TranslateEdge<I, E> translateEdge;

  /** IPC Client */

  private final WorkerClient<I, V, E> workerClient;

  /** IPC Server */

  private final WorkerServer<I, V, E> workerServer;

  /** Request processor for aggregator requests */

  private final WorkerAggregatorRequestProcessor

  workerAggregatorRequestProcessor;

  /** Master info */

  private MasterInfo masterInfo = new MasterInfo();

  /** List of workers */

  private List<WorkerInfo> workerInfoList = Lists.newArrayList();

  /** Have the partition exchange children (workers) changed? */

  private final BspEvent partitionExchangeChildrenChanged;

  /** Worker Context */

  private final WorkerContext workerContext;

  /** Handler for aggregators */

  private final WorkerAggregatorHandler globalCommHandler;

  /** Superstep output */

  private final SuperstepOutput<I, V, E> superstepOutput;

  /** array of observers to call back to */

  private final WorkerObserver[] observers;

  /** Writer for worker progress */

  private final WorkerProgressWriter workerProgressWriter;

  // Per-Superstep Metrics

  /** Timer for WorkerContext#postSuperstep */

  private GiraphTimer wcPostSuperstepTimer;

  /** Time spent waiting on requests to finish */

  private GiraphTimer waitRequestsTimer;

  /**

   * Constructor for setting up the worker.

   *

   * @param context Mapper context

   * @param graphTaskManager GraphTaskManager for this compute node

   * @throws IOException

   * @throws InterruptedException

   */

  public BspServiceWorker(

    Mapper<?, ?, ?, ?>.Context context,

    GraphTaskManager<I, V, E> graphTaskManager)

    throws IOException, InterruptedException {

    super(context, graphTaskManager);

    ImmutableClassesGiraphConfiguration<I, V, E> conf = getConfiguration();

    localData = new LocalData<>(conf);

    translateEdge = getConfiguration().edgeTranslationInstance();

    if (translateEdge != null) {

      translateEdge.initialize(this);

    }

    partitionExchangeChildrenChanged = new PredicateLock(context);

    registerBspEvent(partitionExchangeChildrenChanged);

    workerGraphPartitioner =

        getGraphPartitionerFactory().createWorkerGraphPartitioner();

    workerInfo = new WorkerInfo();

    workerServer = new NettyWorkerServer<I, V, E>(conf, this, context,

        graphTaskManager.createUncaughtExceptionHandler());

    workerInfo.setInetSocketAddress(workerServer.getMyAddress());

    workerInfo.setTaskId(getTaskPartition());

    workerClient = new NettyWorkerClient<I, V, E>(context, conf, this,

        graphTaskManager.createUncaughtExceptionHandler());

    workerAggregatorRequestProcessor =

        new NettyWorkerAggregatorRequestProcessor(getContext(), conf, this);

    globalCommHandler = new WorkerAggregatorHandler(this, conf, context);

    workerContext = conf.createWorkerContext();

    workerContext.setWorkerGlobalCommUsage(globalCommHandler);

    superstepOutput = conf.createSuperstepOutput(context);

    if (conf.isJMapHistogramDumpEnabled()) {

      conf.addWorkerObserverClass(JMapHistoDumper.class);

    }

    if (conf.isReactiveJmapHistogramDumpEnabled()) {

      conf.addWorkerObserverClass(ReactiveJMapHistoDumper.class);

    }

    observers = conf.createWorkerObservers();

    WorkerProgress.get().setTaskId(getTaskPartition());

    workerProgressWriter = conf.trackJobProgressOnClient() ?

        new WorkerProgressWriter(graphTaskManager.getJobProgressTracker()) :

        null;

    GiraphMetrics.get().addSuperstepResetObserver(this);

  }

  @Override

  public void newSuperstep(SuperstepMetricsRegistry superstepMetrics) {

    waitRequestsTimer = new GiraphTimer(superstepMetrics,

        TIMER_WAIT_REQUESTS, TimeUnit.MICROSECONDS);

    wcPostSuperstepTimer = new GiraphTimer(superstepMetrics,

        "worker-context-post-superstep", TimeUnit.MICROSECONDS);

  }

  @Override

  public WorkerContext getWorkerContext() {

    return workerContext;

  }

  @Override

  public WorkerObserver[] getWorkerObservers() {

    return observers;

  }

  @Override

  public WorkerClient<I, V, E> getWorkerClient() {

    return workerClient;

  }

  public LocalData<I, V, E, ? extends Writable> getLocalData() {

    return localData;

  }

  public TranslateEdge<I, E> getTranslateEdge() {

    return translateEdge;

  }

  /**

   * Intended to check the health of the node.  For instance, can it ssh,

   * dmesg, etc. For now, does nothing.

   * TODO: Make this check configurable by the user (i.e. search dmesg for

   * problems).

   *

   * @return True if healthy (always in this case).

   */

  public boolean isHealthy() {

    return true;

  }

  /**

   * Load the vertices/edges from input slits. Do this until all the

   * InputSplits have been processed.

   * All workers will try to do as many InputSplits as they can.  The master

   * will monitor progress and stop this once all the InputSplits have been

   * loaded and check-pointed.  Keep track of the last input split path to

   * ensure the input split cache is flushed prior to marking the last input

   * split complete.

   *

   * Use one or more threads to do the loading.

   *

   * @param inputSplitPathList List of input split paths

   * @param inputSplitsCallableFactory Factory for {@link InputSplitsCallable}s

   * @return Statistics of the vertices and edges loaded

   * @throws InterruptedException

   * @throws KeeperException

   */

  private VertexEdgeCount loadInputSplits(

      List<String> inputSplitPathList,

      CallableFactory<VertexEdgeCount> inputSplitsCallableFactory)

    throws KeeperException, InterruptedException {

    VertexEdgeCount vertexEdgeCount = new VertexEdgeCount();

    // Determine how many threads to use based on the number of input splits

    int maxInputSplitThreads = (inputSplitPathList.size() - 1) /

        getConfiguration().getMaxWorkers() + 1;

    int numThreads = Math.min(getConfiguration().getNumInputSplitsThreads(),

        maxInputSplitThreads);

    if (LOG.isInfoEnabled()) {

      LOG.info("loadInputSplits: Using " + numThreads + " thread(s), " +

          "originally " + getConfiguration().getNumInputSplitsThreads() +

          " threads(s) for " + inputSplitPathList.size() + " total splits.");

    }

    List<VertexEdgeCount> results =

        ProgressableUtils.getResultsWithNCallables(inputSplitsCallableFactory,

            numThreads, "load-%d", getContext());

    for (VertexEdgeCount result : results) {

      vertexEdgeCount = vertexEdgeCount.incrVertexEdgeCount(result);

    }

    workerClient.waitAllRequests();

    return vertexEdgeCount;

  }

  /**

   * Load the mapping entries from the user-defined

   * {@link org.apache.giraph.io.MappingReader}

   *

   * @return Count of mapping entries loaded

   */

  private long loadMapping() throws KeeperException,

    InterruptedException {

    List<String> inputSplitPathList =

        getZkExt().getChildrenExt(mappingInputSplitsPaths.getPath(),

        false, false, true);

    InputSplitPathOrganizer splitOrganizer =

        new InputSplitPathOrganizer(getZkExt(),

            inputSplitPathList, getWorkerInfo().getHostname(),

            getConfiguration().useInputSplitLocality());

    MappingInputSplitsCallableFactory<I, V, E, ? extends Writable>

        mappingInputSplitsCallableFactory =

        new MappingInputSplitsCallableFactory<>(

            getConfiguration().createWrappedMappingInputFormat(),

            splitOrganizer,

            getContext(),

            getConfiguration(),

            this,

            getZkExt());

    long entriesLoaded = 0;

    // Determine how many threads to use based on the number of input splits

    int maxInputSplitThreads = inputSplitPathList.size();

    int numThreads = Math.min(getConfiguration().getNumInputSplitsThreads(),

        maxInputSplitThreads);

    if (LOG.isInfoEnabled()) {

      LOG.info("loadInputSplits: Using " + numThreads + " thread(s), " +

          "originally " + getConfiguration().getNumInputSplitsThreads() +

          " threads(s) for " + inputSplitPathList.size() + " total splits.");

    }

    List<Integer> results =

        ProgressableUtils.getResultsWithNCallables(

            mappingInputSplitsCallableFactory,

            numThreads, "load-mapping-%d", getContext());

    for (Integer result : results) {

      entriesLoaded += result;

    }

    // after all threads finish loading - call postFilling

    localData.getMappingStore().postFilling();

    return entriesLoaded;

  }

  /**

   * Load the vertices from the user-defined

   * {@link org.apache.giraph.io.VertexReader}

   *

   * @return Count of vertices and edges loaded

   */

  private VertexEdgeCount loadVertices() throws KeeperException,

      InterruptedException {

    List<String> inputSplitPathList =

        getZkExt().getChildrenExt(vertexInputSplitsPaths.getPath(),

            false, false, true);

    InputSplitPathOrganizer splitOrganizer =

        new InputSplitPathOrganizer(getZkExt(),

            inputSplitPathList, getWorkerInfo().getHostname(),

            getConfiguration().useInputSplitLocality());

    InputSplitsHandler splitsHandler = new InputSplitsHandler(

        splitOrganizer,

        getZkExt(),

        getContext(),

        BspService.VERTEX_INPUT_SPLIT_RESERVED_NODE,

        BspService.VERTEX_INPUT_SPLIT_FINISHED_NODE);

    VertexInputSplitsCallableFactory<I, V, E> inputSplitsCallableFactory =

        new VertexInputSplitsCallableFactory<I, V, E>(

            getConfiguration().createWrappedVertexInputFormat(),

            getContext(),

            getConfiguration(),

            this,

            splitsHandler,

            getZkExt());

    return loadInputSplits(inputSplitPathList, inputSplitsCallableFactory);

  }

  /**

   * Load the edges from the user-defined

   * {@link org.apache.giraph.io.EdgeReader}.

   *

   * @return Number of edges loaded

   */

  private long loadEdges() throws KeeperException, InterruptedException {

    List<String> inputSplitPathList =

        getZkExt().getChildrenExt(edgeInputSplitsPaths.getPath(),

            false, false, true);

    InputSplitPathOrganizer splitOrganizer =

        new InputSplitPathOrganizer(getZkExt(),

            inputSplitPathList, getWorkerInfo().getHostname(),

            getConfiguration().useInputSplitLocality());

    InputSplitsHandler splitsHandler = new InputSplitsHandler(

        splitOrganizer,

        getZkExt(),

        getContext(),

        BspService.EDGE_INPUT_SPLIT_RESERVED_NODE,

        BspService.EDGE_INPUT_SPLIT_FINISHED_NODE);

    EdgeInputSplitsCallableFactory<I, V, E> inputSplitsCallableFactory =

        new EdgeInputSplitsCallableFactory<I, V, E>(

            getConfiguration().createWrappedEdgeInputFormat(),

            getContext(),

            getConfiguration(),

            this,

            splitsHandler,

            getZkExt());

    return loadInputSplits(inputSplitPathList, inputSplitsCallableFactory).

        getEdgeCount();

  }

  @Override

  public MasterInfo getMasterInfo() {

    return masterInfo;

  }

  @Override

  public List<WorkerInfo> getWorkerInfoList() {

    return workerInfoList;

  }

  /**

   * Ensure the input splits are ready for processing

   *

   * @param inputSplitPaths Input split paths

   * @param inputSplitEvents Input split events

   */

  private void ensureInputSplitsReady(InputSplitPaths inputSplitPaths,

                                      InputSplitEvents inputSplitEvents) {

    while (true) {

      Stat inputSplitsReadyStat;

      try {

        inputSplitsReadyStat = getZkExt().exists(

            inputSplitPaths.getAllReadyPath(), true);

      } catch (KeeperException e) {

        throw new IllegalStateException("ensureInputSplitsReady: " +

            "KeeperException waiting on input splits", e);

      } catch (InterruptedException e) {

        throw new IllegalStateException("ensureInputSplitsReady: " +

            "InterruptedException waiting on input splits", e);

      }

      if (inputSplitsReadyStat != null) {

        break;

      }

      inputSplitEvents.getAllReadyChanged().waitForever();

      inputSplitEvents.getAllReadyChanged().reset();

    }

  }

  /**

   * Mark current worker as done and then wait for all workers

   * to finish processing input splits.

   *

   * @param inputSplitPaths Input split paths

   * @param inputSplitEvents Input split events

   */

  private void markCurrentWorkerDoneThenWaitForOthers(

    InputSplitPaths inputSplitPaths,

    InputSplitEvents inputSplitEvents) {

    String workerInputSplitsDonePath =

        inputSplitPaths.getDonePath() + "/" +

            getWorkerInfo().getHostnameId();

    try {

      getZkExt().createExt(workerInputSplitsDonePath,

          null,

          Ids.OPEN_ACL_UNSAFE,

          CreateMode.PERSISTENT,

          true);

    } catch (KeeperException e) {

      throw new IllegalStateException(

          "markCurrentWorkerDoneThenWaitForOthers: " +

          "KeeperException creating worker done splits", e);

    } catch (InterruptedException e) {

      throw new IllegalStateException(

          "markCurrentWorkerDoneThenWaitForOthers: " +

          "InterruptedException creating worker done splits", e);

    }

    while (true) {

      Stat inputSplitsDoneStat;

      try {

        inputSplitsDoneStat =

            getZkExt().exists(inputSplitPaths.getAllDonePath(),

                true);

      } catch (KeeperException e) {

        throw new IllegalStateException(

            "markCurrentWorkerDoneThenWaitForOthers: " +

            "KeeperException waiting on worker done splits", e);

      } catch (InterruptedException e) {

        throw new IllegalStateException(

            "markCurrentWorkerDoneThenWaitForOthers: " +

            "InterruptedException waiting on worker done splits", e);

      }

      if (inputSplitsDoneStat != null) {

        break;

      }

      inputSplitEvents.getAllDoneChanged().waitForever();

      inputSplitEvents.getAllDoneChanged().reset();

    }

  }

  @Override

  public FinishedSuperstepStats setup() {

    // Unless doing a restart, prepare for computation:

    // 1. Start superstep INPUT_SUPERSTEP (no computation)

    // 2. Wait until the INPUT_SPLIT_ALL_READY_PATH node has been created

    // 3. Process input splits until there are no more.

    // 4. Wait until the INPUT_SPLIT_ALL_DONE_PATH node has been created

    // 5. Process any mutations deriving from add edge requests

    // 6. Wait for superstep INPUT_SUPERSTEP to complete.

    if (getRestartedSuperstep() != UNSET_SUPERSTEP) {

      setCachedSuperstep(getRestartedSuperstep());

      return new FinishedSuperstepStats(0, false, 0, 0, true,

          CheckpointStatus.NONE);

    }

    JSONObject jobState = getJobState();

    if (jobState != null) {

      try {

        if ((ApplicationState.valueOf(jobState.getString(JSONOBJ_STATE_KEY)) ==

            ApplicationState.START_SUPERSTEP) &&

            jobState.getLong(JSONOBJ_SUPERSTEP_KEY) ==

            getSuperstep()) {

          if (LOG.isInfoEnabled()) {

            LOG.info("setup: Restarting from an automated " +

                "checkpointed superstep " +

                getSuperstep() + ", attempt " +

                getApplicationAttempt());

          }

          setRestartedSuperstep(getSuperstep());

          return new FinishedSuperstepStats(0, false, 0, 0, true,

              CheckpointStatus.NONE);

        }

      } catch (JSONException e) {

        throw new RuntimeException(

            "setup: Failed to get key-values from " +

                jobState.toString(), e);

      }

    }

    // Add the partitions that this worker owns

    Collection<? extends PartitionOwner> masterSetPartitionOwners =

        startSuperstep();

    workerGraphPartitioner.updatePartitionOwners(

        getWorkerInfo(), masterSetPartitionOwners);

    /*if[HADOOP_NON_SECURE]

      workerClient.setup();

    else[HADOOP_NON_SECURE]*/

    workerClient.setup(getConfiguration().authenticate());

    /*end[HADOOP_NON_SECURE]*/

    // Initialize aggregator at worker side during setup.

    // Do this just before vertex and edge loading.

    globalCommHandler.prepareSuperstep(workerAggregatorRequestProcessor);

    VertexEdgeCount vertexEdgeCount;

    long entriesLoaded;

    if (getConfiguration().hasMappingInputFormat()) {

      // Ensure the mapping InputSplits are ready for processing

      ensureInputSplitsReady(mappingInputSplitsPaths, mappingInputSplitsEvents);

      getContext().progress();

      try {

        entriesLoaded = loadMapping();

        // successfully loaded mapping

        // now initialize graphPartitionerFactory with this data

        getGraphPartitionerFactory().initialize(localData);

      } catch (InterruptedException e) {

        throw new IllegalStateException(

            "setup: loadMapping failed with InterruptedException", e);

      } catch (KeeperException e) {

        throw new IllegalStateException(

            "setup: loadMapping failed with KeeperException", e);

      }

      getContext().progress();

      if (LOG.isInfoEnabled()) {

        LOG.info("setup: Finally loaded a total of " +

            entriesLoaded + " entries from inputSplits");

      }

      // Workers wait for each other to finish, coordinated by master

      markCurrentWorkerDoneThenWaitForOthers(mappingInputSplitsPaths,

          mappingInputSplitsEvents);

      // Print stats for data stored in localData once mapping is fully

      // loaded on all the workers

      localData.printStats();

    }

    if (getConfiguration().hasVertexInputFormat()) {

      // Ensure the vertex InputSplits are ready for processing

      ensureInputSplitsReady(vertexInputSplitsPaths, vertexInputSplitsEvents);

      getContext().progress();

      try {

        vertexEdgeCount = loadVertices();

      } catch (InterruptedException e) {

        throw new IllegalStateException(

            "setup: loadVertices failed with InterruptedException", e);

      } catch (KeeperException e) {

        throw new IllegalStateException(

            "setup: loadVertices failed with KeeperException", e);

      }

      getContext().progress();

    } else {

      vertexEdgeCount = new VertexEdgeCount();

    }

    WorkerProgress.get().finishLoadingVertices();

    if (getConfiguration().hasEdgeInputFormat()) {

      // Ensure the edge InputSplits are ready for processing

      ensureInputSplitsReady(edgeInputSplitsPaths, edgeInputSplitsEvents);

      getContext().progress();

      try {

        vertexEdgeCount = vertexEdgeCount.incrVertexEdgeCount(0, loadEdges());

      } catch (InterruptedException e) {

        throw new IllegalStateException(

            "setup: loadEdges failed with InterruptedException", e);

      } catch (KeeperException e) {

        throw new IllegalStateException(

            "setup: loadEdges failed with KeeperException", e);

      }

      getContext().progress();

    }

    WorkerProgress.get().finishLoadingEdges();

    if (LOG.isInfoEnabled()) {

      LOG.info("setup: Finally loaded a total of " + vertexEdgeCount);

    }

    if (getConfiguration().hasVertexInputFormat()) {

      // Workers wait for each other to finish, coordinated by master

      markCurrentWorkerDoneThenWaitForOthers(vertexInputSplitsPaths,

          vertexInputSplitsEvents);

    }

    if (getConfiguration().hasEdgeInputFormat()) {

      // Workers wait for each other to finish, coordinated by master

      markCurrentWorkerDoneThenWaitForOthers(edgeInputSplitsPaths,

          edgeInputSplitsEvents);

    }

    // Create remaining partitions owned by this worker.

    for (PartitionOwner partitionOwner : masterSetPartitionOwners) {

      if (partitionOwner.getWorkerInfo().equals(getWorkerInfo()) &&

          !getPartitionStore().hasPartition(

              partitionOwner.getPartitionId())) {

        Partition<I, V, E> partition =

            getConfiguration().createPartition(

                partitionOwner.getPartitionId(), getContext());

        getPartitionStore().addPartition(partition);

      }

    }

    // remove mapping store if possible

    localData.removeMappingStoreIfPossible();

    if (getConfiguration().hasEdgeInputFormat()) {

      // Move edges from temporary storage to their source vertices.

      getServerData().getEdgeStore().moveEdgesToVertices();

    }

    // Generate the partition stats for the input superstep and process

    // if necessary

    List<PartitionStats> partitionStatsList =

        new ArrayList<PartitionStats>();

    for (Integer partitionId : getPartitionStore().getPartitionIds()) {

      Partition<I, V, E> partition =

          getPartitionStore().getOrCreatePartition(partitionId);

      PartitionStats partitionStats =

          new PartitionStats(partition.getId(),

              partition.getVertexCount(),

              0,

              partition.getEdgeCount(),

              0, 0);

      partitionStatsList.add(partitionStats);

      getPartitionStore().putPartition(partition);

    }

    workerGraphPartitioner.finalizePartitionStats(

        partitionStatsList, getPartitionStore());

    return finishSuperstep(partitionStatsList, null);

  }

  /**

   * Register the health of this worker for a given superstep

   *

   * @param superstep Superstep to register health on

   */

  private void registerHealth(long superstep) {

    JSONArray hostnamePort = new JSONArray();

    hostnamePort.put(getHostname());

    hostnamePort.put(workerInfo.getPort());

    String myHealthPath = null;

    if (isHealthy()) {

      myHealthPath = getWorkerInfoHealthyPath(getApplicationAttempt(),

          getSuperstep());

    } else {

      myHealthPath = getWorkerInfoUnhealthyPath(getApplicationAttempt(),

          getSuperstep());

    }

    myHealthPath = myHealthPath + "/" + workerInfo.getHostnameId();

    try {

      myHealthZnode = getZkExt().createExt(

          myHealthPath,

          WritableUtils.writeToByteArray(workerInfo),

          Ids.OPEN_ACL_UNSAFE,

          CreateMode.EPHEMERAL,

          true);

    } catch (KeeperException.NodeExistsException e) {

      LOG.warn("registerHealth: myHealthPath already exists (likely " +

          "from previous failure): " + myHealthPath +

          ".  Waiting for change in attempts " +

          "to re-join the application");

      getApplicationAttemptChangedEvent().waitForever();

      if (LOG.isInfoEnabled()) {

        LOG.info("registerHealth: Got application " +

            "attempt changed event, killing self");

      }

      throw new IllegalStateException(

          "registerHealth: Trying " +

              "to get the new application attempt by killing self", e);

    } catch (KeeperException e) {

      throw new IllegalStateException("Creating " + myHealthPath +

          " failed with KeeperException", e);

    } catch (InterruptedException e) {

      throw new IllegalStateException("Creating " + myHealthPath +

          " failed with InterruptedException", e);

    }

    if (LOG.isInfoEnabled()) {

      LOG.info("registerHealth: Created my health node for attempt=" +

          getApplicationAttempt() + ", superstep=" +

          getSuperstep() + " with " + myHealthZnode +

          " and workerInfo= " + workerInfo);

    }

  }

  /**

   * Do this to help notify the master quicker that this worker has failed.

   */

  private void unregisterHealth() {

    LOG.error("unregisterHealth: Got failure, unregistering health on " +

        myHealthZnode + " on superstep " + getSuperstep());

    try {

      getZkExt().deleteExt(myHealthZnode, -1, false);

    } catch (InterruptedException e) {

      throw new IllegalStateException(

          "unregisterHealth: InterruptedException - Couldn't delete " +

              myHealthZnode, e);

    } catch (KeeperException e) {

      throw new IllegalStateException(

          "unregisterHealth: KeeperException - Couldn't delete " +

              myHealthZnode, e);

    }

  }

  @Override

  public void failureCleanup() {

    unregisterHealth();

  }

  @Override

  public Collection<? extends PartitionOwner> startSuperstep() {

    // Algorithm:

    // 1. Communication service will combine message from previous

    //    superstep

    // 2. Register my health for the next superstep.

    // 3. Wait until the partition assignment is complete and get it

    // 4. Get the aggregator values from the previous superstep

    if (getSuperstep() != INPUT_SUPERSTEP) {

      workerServer.prepareSuperstep();

    }

    registerHealth(getSuperstep());

    String addressesAndPartitionsPath =

        getAddressesAndPartitionsPath(getApplicationAttempt(),

            getSuperstep());

    AddressesAndPartitionsWritable addressesAndPartitions =

        new AddressesAndPartitionsWritable(

            workerGraphPartitioner.createPartitionOwner().getClass());

    try {

      while (getZkExt().exists(addressesAndPartitionsPath, true) ==

          null) {

        getAddressesAndPartitionsReadyChangedEvent().waitForever();

        getAddressesAndPartitionsReadyChangedEvent().reset();

      }

      WritableUtils.readFieldsFromZnode(

          getZkExt(),

          addressesAndPartitionsPath,

          false,

          null,

          addressesAndPartitions);

    } catch (KeeperException e) {

      throw new IllegalStateException(

          "startSuperstep: KeeperException getting assignments", e);

    } catch (InterruptedException e) {

      throw new IllegalStateException(

          "startSuperstep: InterruptedException getting assignments", e);

    }

    workerInfoList.clear();

    workerInfoList = addressesAndPartitions.getWorkerInfos();

    masterInfo = addressesAndPartitions.getMasterInfo();

    if (LOG.isInfoEnabled()) {

      LOG.info("startSuperstep: " + masterInfo);

      LOG.info("startSuperstep: Ready for computation on superstep " +

          getSuperstep() + " since worker " +

          "selection and vertex range assignments are done in " +

          addressesAndPartitionsPath);

    }

    getContext().setStatus("startSuperstep: " +

        getGraphTaskManager().getGraphFunctions().toString() +

        " - Attempt=" + getApplicationAttempt() +

        ", Superstep=" + getSuperstep());

    if (LOG.isDebugEnabled()) {

      LOG.debug("startSuperstep: addressesAndPartitions" +

          addressesAndPartitions.getWorkerInfos());

      for (PartitionOwner partitionOwner : addressesAndPartitions

          .getPartitionOwners()) {

        LOG.debug(partitionOwner.getPartitionId() + " " +

            partitionOwner.getWorkerInfo());

      }

    }

    return addressesAndPartitions.getPartitionOwners();

  }

  @Override

  public FinishedSuperstepStats finishSuperstep(

      List<PartitionStats> partitionStatsList,

      GiraphTimerContext superstepTimerContext) {

    // This barrier blocks until success (or the master signals it to

    // restart).

    //

    // Master will coordinate the barriers and aggregate "doneness" of all

    // the vertices.  Each worker will:

    // 1. Ensure that the requests are complete

    // 2. Execute user postSuperstep() if necessary.

    // 3. Save aggregator values that are in use.

    // 4. Report the statistics (vertices, edges, messages, etc.)

    //    of this worker

    // 5. Let the master know it is finished.

    // 6. Wait for the master's superstep info, and check if done

    waitForRequestsToFinish();

    getGraphTaskManager().notifyFinishedCommunication();

    long workerSentMessages = 0;

    long workerSentMessageBytes = 0;

    long localVertices = 0;

    for (PartitionStats partitionStats : partitionStatsList) {

      workerSentMessages += partitionStats.getMessagesSentCount();

      workerSentMessageBytes += partitionStats.getMessageBytesSentCount();

      localVertices += partitionStats.getVertexCount();

    }

    if (getSuperstep() != INPUT_SUPERSTEP) {

      postSuperstepCallbacks();

    }

    globalCommHandler.finishSuperstep(workerAggregatorRequestProcessor);

    MessageStore<I, Writable> incomingMessageStore =

        getServerData().getIncomingMessageStore();

    if (incomingMessageStore instanceof AsyncMessageStoreWrapper) {

      ((AsyncMessageStoreWrapper) incomingMessageStore).waitToComplete();

    }

    if (LOG.isInfoEnabled()) {

      LOG.info("finishSuperstep: Superstep " + getSuperstep() +

          ", messages = " + workerSentMessages + " " +

          ", message bytes = " + workerSentMessageBytes + " , " +

          MemoryUtils.getRuntimeMemoryStats());

    }

    if (superstepTimerContext != null) {

      superstepTimerContext.stop();

    }

    writeFinshedSuperstepInfoToZK(partitionStatsList,

      workerSentMessages, workerSentMessageBytes);

    LoggerUtils.setStatusAndLog(getContext(), LOG, Level.INFO,

        "finishSuperstep: (waiting for rest " +

            "of workers) " +

            getGraphTaskManager().getGraphFunctions().toString() +

            " - Attempt=" + getApplicationAttempt() +

            ", Superstep=" + getSuperstep());

    String superstepFinishedNode =

        getSuperstepFinishedPath(getApplicationAttempt(), getSuperstep());

    waitForOtherWorkers(superstepFinishedNode);

    GlobalStats globalStats = new GlobalStats();

    SuperstepClasses superstepClasses = new SuperstepClasses();

    WritableUtils.readFieldsFromZnode(

        getZkExt(), superstepFinishedNode, false, null, globalStats,

        superstepClasses);

    if (LOG.isInfoEnabled()) {

      LOG.info("finishSuperstep: Completed superstep " + getSuperstep() +

          " with global stats " + globalStats + " and classes " +

          superstepClasses);

    }

    getContext().setStatus("finishSuperstep: (all workers done) " +

        getGraphTaskManager().getGraphFunctions().toString() +

        " - Attempt=" + getApplicationAttempt() +

        ", Superstep=" + getSuperstep());

    incrCachedSuperstep();

    getConfiguration().updateSuperstepClasses(superstepClasses);

    return new FinishedSuperstepStats(

        localVertices,

        globalStats.getHaltComputation(),

        globalStats.getVertexCount(),

        globalStats.getEdgeCount(),

        false,

        globalStats.getCheckpointStatus());

  }

  /**

   * Handle post-superstep callbacks

   */

  private void postSuperstepCallbacks() {

    GiraphTimerContext timerContext = wcPostSuperstepTimer.time();

    getWorkerContext().postSuperstep();

    timerContext.stop();

    getContext().progress();

    for (WorkerObserver obs : getWorkerObservers()) {

      obs.postSuperstep(getSuperstep());

      getContext().progress();

    }

  }

  /**

   * Wait for all the requests to finish.

   */

  private void waitForRequestsToFinish() {

    if (LOG.isInfoEnabled()) {

      LOG.info("finishSuperstep: Waiting on all requests, superstep " +

          getSuperstep() + " " +

          MemoryUtils.getRuntimeMemoryStats());

    }

    GiraphTimerContext timerContext = waitRequestsTimer.time();

    workerClient.waitAllRequests();

    timerContext.stop();

  }

  /**

   * Wait for all the other Workers to finish the superstep.

   *

   * @param superstepFinishedNode ZooKeeper path to wait on.

   */

  private void waitForOtherWorkers(String superstepFinishedNode) {

    try {

      while (getZkExt().exists(superstepFinishedNode, true) == null) {

        getSuperstepFinishedEvent().waitForever();

        getSuperstepFinishedEvent().reset();

      }

    } catch (KeeperException e) {

      throw new IllegalStateException(

          "finishSuperstep: Failed while waiting for master to " +

              "signal completion of superstep " + getSuperstep(), e);

    } catch (InterruptedException e) {

      throw new IllegalStateException(

          "finishSuperstep: Failed while waiting for master to " +

              "signal completion of superstep " + getSuperstep(), e);

    }

  }

  /**

   * Write finished superstep info to ZooKeeper.

   *

   * @param partitionStatsList List of partition stats from superstep.

   * @param workerSentMessages Number of messages sent in superstep.

   * @param workerSentMessageBytes Number of message bytes sent

   *                               in superstep.

   */

  private void writeFinshedSuperstepInfoToZK(

      List<PartitionStats> partitionStatsList, long workerSentMessages,

      long workerSentMessageBytes) {

    Collection<PartitionStats> finalizedPartitionStats =

        workerGraphPartitioner.finalizePartitionStats(

            partitionStatsList, getPartitionStore());

    List<PartitionStats> finalizedPartitionStatsList =

        new ArrayList<PartitionStats>(finalizedPartitionStats);

    byte[] partitionStatsBytes =

        WritableUtils.writeListToByteArray(finalizedPartitionStatsList);

    WorkerSuperstepMetrics metrics = new WorkerSuperstepMetrics();

    metrics.readFromRegistry();

    byte[] metricsBytes = WritableUtils.writeToByteArray(metrics);

    JSONObject workerFinishedInfoObj = new JSONObject();

    try {

      workerFinishedInfoObj.put(JSONOBJ_PARTITION_STATS_KEY,

          Base64.encodeBytes(partitionStatsBytes));

      workerFinishedInfoObj.put(JSONOBJ_NUM_MESSAGES_KEY, workerSentMessages);

      workerFinishedInfoObj.put(JSONOBJ_NUM_MESSAGE_BYTES_KEY,

        workerSentMessageBytes);

      workerFinishedInfoObj.put(JSONOBJ_METRICS_KEY,

          Base64.encodeBytes(metricsBytes));

    } catch (JSONException e) {

      throw new RuntimeException(e);

    }

    String finishedWorkerPath =

        getWorkerFinishedPath(getApplicationAttempt(), getSuperstep()) +

        "/" + getHostnamePartitionId();

    try {

      getZkExt().createExt(finishedWorkerPath,

          workerFinishedInfoObj.toString().getBytes(Charset.defaultCharset()),

          Ids.OPEN_ACL_UNSAFE,

          CreateMode.PERSISTENT,

          true);

    } catch (KeeperException.NodeExistsException e) {

      LOG.warn("finishSuperstep: finished worker path " +

          finishedWorkerPath + " already exists!");

    } catch (KeeperException e) {

      throw new IllegalStateException("Creating " + finishedWorkerPath +

          " failed with KeeperException", e);

    } catch (InterruptedException e) {

      throw new IllegalStateException("Creating " + finishedWorkerPath +

          " failed with InterruptedException", e);

    }

  }

  /**

   * Save the vertices using the user-defined VertexOutputFormat from our

   * vertexArray based on the split.

   *

   * @param numLocalVertices Number of local vertices

   * @throws InterruptedException

   */

  private void saveVertices(long numLocalVertices) throws IOException,

      InterruptedException {

    ImmutableClassesGiraphConfiguration<I, V, E>  conf = getConfiguration();

    if (conf.getVertexOutputFormatClass() == null) {

      LOG.warn("saveVertices: " +

          GiraphConstants.VERTEX_OUTPUT_FORMAT_CLASS +

          " not specified -- there will be no saved output");

      return;

    }

    if (conf.doOutputDuringComputation()) {

      if (LOG.isInfoEnabled()) {

        LOG.info("saveVertices: The option for doing output during " +

            "computation is selected, so there will be no saving of the " +

            "output in the end of application");

      }

      return;

    }

    final int numPartitions = getPartitionStore().getNumPartitions();

    int numThreads = Math.min(getConfiguration().getNumOutputThreads(),

        numPartitions);

    LoggerUtils.setStatusAndLog(getContext(), LOG, Level.INFO,

        "saveVertices: Starting to save " + numLocalVertices + " vertices " +

            "using " + numThreads + " threads");

    final VertexOutputFormat<I, V, E> vertexOutputFormat =

        getConfiguration().createWrappedVertexOutputFormat();

    final Queue<Integer> partitionIdQueue =

        (numPartitions == 0) ? new LinkedList<Integer>() :

            new ArrayBlockingQueue<Integer>(numPartitions);

    Iterables.addAll(partitionIdQueue, getPartitionStore().getPartitionIds());

    long verticesToStore = 0;

    PartitionStore<I, V, E> partitionStore = getPartitionStore();

    for (int partitionId : partitionStore.getPartitionIds()) {

      Partition<I, V, E> partition =

        partitionStore.getOrCreatePartition(partitionId);

      verticesToStore += partition.getVertexCount();

      partitionStore.putPartition(partition);

    }

    WorkerProgress.get().startStoring(

        verticesToStore, getPartitionStore().getNumPartitions());

    CallableFactory<Void> callableFactory = new CallableFactory<Void>() {

      @Override

      public Callable<Void> newCallable(int callableId) {

        return new Callable<Void>() {

          /** How often to update WorkerProgress */

          private static final long VERTICES_TO_UPDATE_PROGRESS = 100000;

          @Override

          public Void call() throws Exception {

            VertexWriter<I, V, E> vertexWriter =

                vertexOutputFormat.createVertexWriter(getContext());

            vertexWriter.setConf(getConfiguration());

            vertexWriter.initialize(getContext());

            long nextPrintVertices = 0;

            long nextUpdateProgressVertices = VERTICES_TO_UPDATE_PROGRESS;

            long nextPrintMsecs = System.currentTimeMillis() + 15000;

            int partitionIndex = 0;

            int numPartitions = getPartitionStore().getNumPartitions();

            while (!partitionIdQueue.isEmpty()) {

              Integer partitionId = partitionIdQueue.poll();

              if (partitionId == null) {

                break;

              }

              Partition<I, V, E> partition =

                  getPartitionStore().getOrCreatePartition(partitionId);

              long verticesWritten = 0;

              for (Vertex<I, V, E> vertex : partition) {

                vertexWriter.writeVertex(vertex);

                ++verticesWritten;

                // Update status at most every 250k vertices or 15 seconds

                if (verticesWritten > nextPrintVertices &&

                    System.currentTimeMillis() > nextPrintMsecs) {

                  LoggerUtils.setStatusAndLog(getContext(), LOG, Level.INFO,

                      "saveVertices: Saved " + verticesWritten + " out of " +

                          partition.getVertexCount() + " partition vertices, " +

                          "on partition " + partitionIndex +

                          " out of " + numPartitions);

                  nextPrintMsecs = System.currentTimeMillis() + 15000;

                  nextPrintVertices = verticesWritten + 250000;

                }

                if (verticesWritten >= nextUpdateProgressVertices) {

                  WorkerProgress.get().addVerticesStored(

                      VERTICES_TO_UPDATE_PROGRESS);

                  nextUpdateProgressVertices += VERTICES_TO_UPDATE_PROGRESS;

                }

              }

              getPartitionStore().putPartition(partition);

              ++partitionIndex;

              WorkerProgress.get().addVerticesStored(

                  verticesWritten % VERTICES_TO_UPDATE_PROGRESS);

              WorkerProgress.get().incrementPartitionsStored();

            }

            vertexWriter.close(getContext()); // the temp results are saved now

            return null;