/*

  * 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.block_app.framework.output;

 import java.util.HashMap;

 import java.util.Map;

 import java.util.Queue;

 import java.util.concurrent.Callable;

 import java.util.concurrent.ConcurrentLinkedQueue;

 import org.apache.giraph.block_app.framework.api.BlockOutputApi;

 import org.apache.giraph.conf.GiraphConstants;

 import org.apache.giraph.utils.CallableFactory;

 import org.apache.giraph.utils.ProgressableUtils;

 import org.apache.hadoop.conf.Configuration;

 import org.apache.hadoop.util.Progressable;

 /**

  * Handler for blocks output - keeps track of outputs and writers created

  */

 @SuppressWarnings("unchecked")

 public class BlockOutputHandle implements BlockOutputApi {

   private transient Configuration conf;

   private transient Progressable progressable;

   private final Map<String, BlockOutputDesc> outputDescMap;

   private final Map<String, Queue<BlockOutputWriter>> freeWriters =

       new HashMap<>();

   private final Map<String, Queue<BlockOutputWriter>> occupiedWriters =

       new HashMap<>();

   public BlockOutputHandle() {

     outputDescMap = null;

   }

   public BlockOutputHandle(String jobIdentifier, Configuration conf,

       Progressable hadoopProgressable) {

     outputDescMap = BlockOutputFormat.createInitAndCheckOutputDescsMap(

         conf, jobIdentifier);

     for (Map.Entry<String, BlockOutputDesc> entry : outputDescMap.entrySet()) {

       entry.getValue().preWriting();

       freeWriters.put(entry.getKey(),

           new ConcurrentLinkedQueue<BlockOutputWriter>());

       occupiedWriters.put(entry.getKey(),

           new ConcurrentLinkedQueue<BlockOutputWriter>());

     }

     initialize(conf, hadoopProgressable);

   }

   public void initialize(Configuration conf, Progressable progressable) {

     this.conf = conf;

     this.progressable = progressable;

   }

   @Override

   public <OW extends BlockOutputWriter, OD extends BlockOutputDesc<OW>>

   OD getOutputDesc(String confOption) {

     if (outputDescMap == null) {

       throw new IllegalArgumentException(

           "Output cannot be used with checkpointing");

     }

     return (OD) outputDescMap.get(confOption);

   }

   @Override

   public <OW extends BlockOutputWriter> OW getWriter(String confOption) {

     if (outputDescMap == null) {

       throw new IllegalArgumentException(

           "Output cannot be used with checkpointing");

     }

     OW outputWriter = (OW) freeWriters.get(confOption).poll();

     if (outputWriter == null) {

       outputWriter = (OW) outputDescMap.get(confOption).createOutputWriter(

           conf, progressable);

     }

     occupiedWriters.get(confOption).add(outputWriter);

     return outputWriter;

   }

   public void returnAllWriters() {

     for (Map.Entry<String, Queue<BlockOutputWriter>> entry :

         occupiedWriters.entrySet()) {

       freeWriters.get(entry.getKey()).addAll(entry.getValue());

       entry.getValue().clear();

     }

   }

   public void closeAllWriters() {

     final Queue<BlockOutputWriter> allWriters = new ConcurrentLinkedQueue<>();

     for (Queue<BlockOutputWriter> blockOutputWriters : freeWriters.values()) {

       allWriters.addAll(blockOutputWriters);

     }

     if (allWriters.isEmpty()) {

       return;

     }

     // Closing writers can take time - use multiple threads and call progress

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

       @Override

       public Callable<Void> newCallable(int callableId) {

         return new Callable<Void>() {

           @Override

           public Void call() throws Exception {

             BlockOutputWriter writer = allWriters.poll();

             while (writer != null) {

               writer.close();

               writer = allWriters.poll();

             }

             return null;

           }

         };

       }

     };

     ProgressableUtils.getResultsWithNCallables(callableFactory,

         Math.min(GiraphConstants.NUM_OUTPUT_THREADS.get(conf),

             allWriters.size()), "close-writers-%d", progressable);

     // Close all output formats

     for (BlockOutputDesc outputDesc : outputDescMap.values()) {

       outputDesc.postWriting();

     }

   }

 }