/*
    * 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.master.input;
  
  import org.apache.giraph.worker.WorkerInfo;
  
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.concurrent.atomic.AtomicInteger;
  
  /**
   * Organizer for mapping splits on master. Mapping splits need all to be
   * given to all workers, unlike vertex and edge splits which are read by
   * exactly one worker each
   */
  public class MappingInputSplitsMasterOrganizer
      implements InputSplitsMasterOrganizer {
    /** List of splits */
    private final List<byte[]> splits;
    /** Map from worker task id to atomic pointer in splits list */
    private final Map<Integer, AtomicInteger>
        workerTaskIdToNextSplitIndexMap;
  
    /**
     * Constructor
     *
     * @param serializedSplits Splits
     * @param workers List of workers
     */
    public MappingInputSplitsMasterOrganizer(List<byte[]> serializedSplits,
        List<WorkerInfo> workers) {
      this.splits = serializedSplits;
      workerTaskIdToNextSplitIndexMap = new HashMap<>();
      for (WorkerInfo worker : workers) {
        workerTaskIdToNextSplitIndexMap.put(
            worker.getTaskId(), new AtomicInteger(0));
      }
    }
  
    @Override
    public byte[] getSerializedSplitFor(int workerTaskId) {
      AtomicInteger nextSplitIndex =
          workerTaskIdToNextSplitIndexMap.get(workerTaskId);
      int splitIndex = nextSplitIndex.getAndIncrement();
      return splitIndex < splits.size() ? splits.get(splitIndex) : null;
    }
  }