/*
    * 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);
     }
   }
 }