giraph-core/src/main/java/org/apache/giraph/ooc/OutOfCoreIOScheduler.java

/*
 * 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.giraph.ooc;

import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;
import org.apache.giraph.conf.IntConfOption;
import org.apache.giraph.ooc.command.IOCommand;
import org.apache.giraph.ooc.command.LoadPartitionIOCommand;
import org.apache.giraph.ooc.command.StoreDataBufferIOCommand;
import org.apache.giraph.ooc.command.StoreIncomingMessageIOCommand;
import org.apache.giraph.ooc.command.StorePartitionIOCommand;
import org.apache.giraph.ooc.command.WaitIOCommand;
import org.apache.giraph.ooc.policy.OutOfCoreOracle;
import org.apache.log4j.Logger;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;

import static com.google.common.base.Preconditions.checkNotNull;

/**
 * Representation of IO thread scheduler for out-of-core mechanism
 */
public class OutOfCoreIOScheduler {
  /**
   * If an IO thread does not have any command to do, it waits for certain a
   * period and check back again to see if there exist any command to perform.
   * This constant determines this wait period in milliseconds.
   */
  public static final IntConfOption OOC_WAIT_INTERVAL =
      new IntConfOption("giraph.oocWaitInterval", 1000,
          "Duration (in milliseconds) which IO threads in out-of-core " +
              "mechanism would wait until a command becomes available");
  /** Class logger. */
  private static final Logger LOG =
      Logger.getLogger(OutOfCoreIOScheduler.class);
  /** Out-of-core engine */
  private final OutOfCoreEngine oocEngine;
  /** How much an IO thread should wait if there is no IO command */
  private final int waitInterval;
  /**
   * Queue of IO commands for loading partitions to memory. Load commands are
   * urgent and should be done once loading data is a viable IO command.
   */
  private final List<Queue<IOCommand>> threadLoadCommandQueue;
  /** Whether IO threads should terminate */
  private volatile boolean shouldTerminate;

  /**
   * Constructor
   *
   * @param conf configuration
   * @param oocEngine out-of-core engine
   * @param numDisks number of disks (IO threads)
   */
  OutOfCoreIOScheduler(final ImmutableClassesGiraphConfiguration conf,
                       OutOfCoreEngine oocEngine, int numDisks) {
    this.oocEngine = oocEngine;
    this.waitInterval = OOC_WAIT_INTERVAL.get(conf);
    threadLoadCommandQueue = new ArrayList<>(numDisks);
    for (int i = 0; i < numDisks; ++i) {
      threadLoadCommandQueue.add(
          new ConcurrentLinkedQueue<IOCommand>());
    }
    shouldTerminate = false;
  }

  /**
   * Generate and return the next appropriate IO command for a given thread
   *
   * @param threadId id of the thread ready to execute the next IO command
   * @return next IO command to be executed by the given thread
   */
  public IOCommand getNextIOCommand(int threadId) {
    if (shouldTerminate) {
      return null;
    }
    IOCommand command = null;
    do {
      if (command != null && LOG.isInfoEnabled()) {
        LOG.info("getNextIOCommand: command " + command + " was proposed to " +
            "the oracle, but got denied. Generating another command!");
      }
      OutOfCoreOracle.IOAction[] actions =
          oocEngine.getOracle().getNextIOActions();
      if (LOG.isDebugEnabled()) {
        LOG.debug("getNextIOCommand: actions are " + Arrays.toString(actions));
      }
      // Check whether there are any urgent outstanding load requests
      if (!threadLoadCommandQueue.get(threadId).isEmpty()) {
        // Check whether loading a partition is a viable (allowed) action to do
        boolean canLoad = false;
        for (OutOfCoreOracle.IOAction action : actions) {
          if (action == OutOfCoreOracle.IOAction.LOAD_PARTITION ||
              action == OutOfCoreOracle.IOAction.LOAD_UNPROCESSED_PARTITION ||
              action == OutOfCoreOracle.IOAction.LOAD_TO_SWAP_PARTITION ||
              action == OutOfCoreOracle.IOAction.URGENT_LOAD_PARTITION) {
            canLoad = true;
            break;
          }
        }
        if (canLoad) {
          command = threadLoadCommandQueue.get(threadId).poll();
          checkNotNull(command);
          if (oocEngine.getOracle().approve(command)) {
            return command;
          } else {
            // Loading is not viable at this moment. We should put the command
            // back in the load queue and wait until loading becomes viable.
            threadLoadCommandQueue.get(threadId).offer(command);
          }
        }
      }
      command = null;
      for (OutOfCoreOracle.IOAction action : actions) {
        Integer partitionId;
        switch (action) {
        case STORE_MESSAGES_AND_BUFFERS:
          partitionId = oocEngine.getMetaPartitionManager()
              .getOffloadPartitionBufferId(threadId);
          if (partitionId != null) {
            command = new StoreDataBufferIOCommand(oocEngine, partitionId,
                StoreDataBufferIOCommand.DataBufferType.PARTITION);
          } else {
            partitionId = oocEngine.getMetaPartitionManager()
                .getOffloadMessageBufferId(threadId);
            if (partitionId != null) {
              command = new StoreDataBufferIOCommand(oocEngine, partitionId,
                  StoreDataBufferIOCommand.DataBufferType.MESSAGE);
            } else {
              partitionId = oocEngine.getMetaPartitionManager()
                  .getOffloadMessageId(threadId);
              if (partitionId != null) {
                command = new StoreIncomingMessageIOCommand(oocEngine,
                    partitionId);
              }
            }
          }
          break;
        case STORE_PROCESSED_PARTITION:
          partitionId = oocEngine.getMetaPartitionManager()
              .getOffloadPartitionId(threadId);
          if (partitionId != null &&
              oocEngine.getMetaPartitionManager()
                  .isPartitionProcessed(partitionId)) {
            command = new StorePartitionIOCommand(oocEngine, partitionId);
          }
          break;
        case STORE_PARTITION:
          partitionId = oocEngine.getMetaPartitionManager()
              .getOffloadPartitionId(threadId);
          if (partitionId != null) {
            command = new StorePartitionIOCommand(oocEngine, partitionId);
          }
          break;
        case LOAD_UNPROCESSED_PARTITION:
          partitionId = oocEngine.getMetaPartitionManager()
              .getLoadPartitionId(threadId);
          if (partitionId != null &&
              !oocEngine.getMetaPartitionManager()
                  .isPartitionProcessed(partitionId)) {
            command = new LoadPartitionIOCommand(oocEngine, partitionId,
                oocEngine.getSuperstep());
          }
          break;
        case LOAD_TO_SWAP_PARTITION:
          partitionId = oocEngine.getMetaPartitionManager()
              .getLoadPartitionId(threadId);
          if (partitionId != null &&
              !oocEngine.getMetaPartitionManager()
                  .isPartitionProcessed(partitionId) &&
              oocEngine.getMetaPartitionManager().hasProcessedOnMemory()) {
            command = new LoadPartitionIOCommand(oocEngine, partitionId,
                oocEngine.getSuperstep());
          }
          break;
        case LOAD_PARTITION:
          partitionId = oocEngine.getMetaPartitionManager()
              .getLoadPartitionId(threadId);
          if (partitionId != null) {
            if (oocEngine.getMetaPartitionManager()
                .isPartitionProcessed(partitionId)) {
              command = new LoadPartitionIOCommand(oocEngine, partitionId,
                  oocEngine.getSuperstep() + 1);
            } else {
              command = new LoadPartitionIOCommand(oocEngine, partitionId,
                  oocEngine.getSuperstep());
            }
          }
          break;
        case URGENT_LOAD_PARTITION:
          // Do nothing
          break;
        default:
          throw new IllegalStateException("getNextIOCommand: the IO action " +
              "is not defined!");
        }
        if (command != null) {
          break;
        }
      }
      if (command == null) {
        command = new WaitIOCommand(oocEngine, waitInterval);
      }
    } while (!oocEngine.getOracle().approve(command));
    return command;
  }

  /**
   * Notify IO scheduler that the IO command is completed
   *
   * @param command completed command
   */
  public void ioCommandCompleted(IOCommand command) {
    oocEngine.ioCommandCompleted(command);
  }

  /**
   * Add an IO command to the scheduling queue of the IO scheduler
   *
   * @param ioCommand IO command to add to the scheduler
   */
  public void addIOCommand(IOCommand ioCommand) {
    if (ioCommand instanceof LoadPartitionIOCommand) {
      int ownerThread = oocEngine.getMetaPartitionManager()
          .getOwnerThreadId(ioCommand.getPartitionId());
      threadLoadCommandQueue.get(ownerThread).offer(ioCommand);
    } else {
      throw new IllegalStateException("addIOCommand: IO command type is not " +
          "supported for addition");
    }
  }

  /**
   * Shutdown/Terminate the IO scheduler, and notify all IO threads to halt
   */
  public void shutdown() {
    shouldTerminate = true;
    if (LOG.isInfoEnabled()) {
      LOG.info("shutdown: OutOfCoreIOScheduler shutting down!");
    }
  }
}