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

import org.apache.giraph.comm.WorkerClient;
import org.apache.giraph.comm.requests.AskForInputSplitRequest;
import org.apache.giraph.io.InputType;

import java.util.EnumMap;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * Requests splits from master and keeps track of them
 */
public class WorkerInputSplitsHandler {
  /** Worker info of this worker */
  private final WorkerInfo workerInfo;
  /** Task id of master */
  private final int masterTaskId;
  /** Worker client, used for communication */
  private final WorkerClient workerClient;
  /** Map with currently available splits received from master */
  private final Map<InputType, BlockingQueue<byte[]>> availableInputSplits;

  /**
   * Constructor
   *
   * @param workerInfo   Worker info of this worker
   * @param masterTaskId Task id of master
   * @param workerClient Worker client, used for communication
   */
  public WorkerInputSplitsHandler(WorkerInfo workerInfo, int masterTaskId,
      WorkerClient workerClient) {
    this.workerInfo = workerInfo;
    this.masterTaskId = masterTaskId;
    this.workerClient = workerClient;
    availableInputSplits = new EnumMap<>(InputType.class);
    for (InputType inputType : InputType.values()) {
      availableInputSplits.put(
          inputType, new LinkedBlockingQueue<byte[]>());
    }
  }

  /**
   * Called when an input split has been received from master, adding it to
   * the map
   *
   * @param splitType            Type of split
   * @param serializedInputSplit Split
   */
  public void receivedInputSplit(InputType splitType,
      byte[] serializedInputSplit) {
    try {
      availableInputSplits.get(splitType).put(serializedInputSplit);
    } catch (InterruptedException e) {
      throw new IllegalStateException("Interrupted", e);
    }
  }

  /**
   * Try to reserve an InputSplit for loading.  While InputSplits exists that
   * are not finished, wait until they are.
   *
   * NOTE: iterations on the InputSplit list only halt for each worker when it
   * has scanned the entire list once and found every split marked RESERVED.
   * When a worker fails, its Ephemeral RESERVED znodes will disappear,
   * allowing other iterating workers to claim it's previously read splits.
   * Only when the last worker left iterating on the list fails can a danger
   * of data loss occur. Since worker failure in INPUT_SUPERSTEP currently
   * causes job failure, this is OK. As the failure model evolves, this
   * behavior might need to change. We could add watches on
   * inputSplitFinishedNodes and stop iterating only when all these nodes
   * have been created.
   *
   * @param splitType Type of split
   * @param isFirstSplit Whether this is the first split input thread reads
   * @return reserved InputSplit or null if no unfinished InputSplits exist
   */
  public byte[] reserveInputSplit(InputType splitType, boolean isFirstSplit) {
    // Send request
    workerClient.sendWritableRequest(masterTaskId,
        new AskForInputSplitRequest(
            splitType, workerInfo.getTaskId(), isFirstSplit));
    try {
      // Wait for some split to become available
      byte[] serializedInputSplit = availableInputSplits.get(splitType).take();
      return serializedInputSplit.length == 0 ? null : serializedInputSplit;
    } catch (InterruptedException e) {
      throw new IllegalStateException("Interrupted", e);
    }
  }
}