/*
    * 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.comm.messages.queue;
  
  import it.unimi.dsi.fastutil.ints.Int2IntArrayMap;
  import it.unimi.dsi.fastutil.ints.Int2IntMap;
  
  import java.io.DataInput;
  import java.io.DataOutput;
  import java.io.IOException;
  import java.util.concurrent.BlockingQueue;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.LinkedBlockingQueue;
  import java.util.concurrent.Semaphore;
  
  import org.apache.giraph.comm.messages.MessageStore;
  import org.apache.giraph.utils.ThreadUtils;
  import org.apache.giraph.utils.VertexIdMessages;
  import org.apache.hadoop.io.Writable;
  import org.apache.hadoop.io.WritableComparable;
  import org.apache.log4j.Logger;
  
  /**
   * This class decouples message receiving and processing
   * into separate threads thus reducing contention.
   * It does not provide message store functionality itself, rather
   * providing a wrapper around existing message stores that
   * can now be used in async mode with only slight modifications.
   * @param <I> Vertex id
   * @param <M> Message data
   */
  public final class AsyncMessageStoreWrapper<I extends WritableComparable,
      M extends Writable> implements MessageStore<I, M> {
  
    /** Logger */
    private static final Logger LOG =
        Logger.getLogger(AsyncMessageStoreWrapper.class);
    /** Pass this id to clear the queues and shutdown all threads
     * started by this processor */
    private static final PartitionMessage SHUTDOWN_QUEUE_MESSAGE =
        new PartitionMessage(-1, null);
    /** Pass this message to clear the queues but keep threads alive */
    private static final PartitionMessage CLEAR_QUEUE_MESSAGE =
        new PartitionMessage(-1, null);
    /** Executor that processes messages in background */
    private static final ExecutorService EXECUTOR_SERVICE =
        Executors.newCachedThreadPool(
            ThreadUtils.createThreadFactory("AsyncMessageStoreWrapper-%d"));
  
    /** Number of threads that will process messages in background */
    private final int threadsCount;
    /** Queue that temporary stores messages */
    private final BlockingQueue<PartitionMessage<I, M>>[] queues;
    /** Map from partition id to thread that process this partition */
    private final Int2IntMap partition2Queue;
    /** Signals that all procesing is done */
    private Semaphore completionSemaphore;
    /** Underlying message store */
    private final MessageStore<I, M> store;
  
    /**
     * Constructs async wrapper around existing message store
     * object. Requires partition list and number of threads
     * to properly initialize background threads and assign partitions.
     * Partitions are assigned to threads in round-robin fashion.
     * It guarantees that all threads have almost the same number of
     * partitions (+-1) no matter how partitions are assigned to this worker.
     * @param store underlying message store to be used in computation
     * @param partitions partitions assigned to this worker
     * @param threadCount number of threads that will be used to process
     *                    messages.
     */
    public AsyncMessageStoreWrapper(MessageStore<I, M> store,
                                    Iterable<Integer> partitions,
                                    int threadCount) {
      this.store = store;
      this.threadsCount = threadCount;
      completionSemaphore = new Semaphore(1 - threadsCount);
      queues = new BlockingQueue[threadsCount];
      partition2Queue = new Int2IntArrayMap();
      LOG.info("AsyncMessageStoreWrapper enabled. Threads= " + threadsCount);
  
      for (int i = 0; i < threadsCount; i++) {
       queues[i] = new LinkedBlockingQueue<>();
       EXECUTOR_SERVICE.submit(new MessageStoreQueueWorker(queues[i]));
     }
 
     int cnt = 0;
     for (int partitionId : partitions) {
       partition2Queue.put(partitionId, cnt++ % threadsCount);
     }
 
   }
 
   @Override
   public boolean isPointerListEncoding() {
     return store.isPointerListEncoding();
   }
 
   @Override
   public Iterable<M> getVertexMessages(I vertexId) {
     return store.getVertexMessages(vertexId);
   }
 
   @Override
   public void clearVertexMessages(I vertexId) {
     store.clearVertexMessages(vertexId);
   }
 
   @Override
   public void clearAll() {
     try {
       for (BlockingQueue<PartitionMessage<I, M>> queue : queues) {
         queue.put(SHUTDOWN_QUEUE_MESSAGE);
       }
       completionSemaphore.acquire();
     } catch (InterruptedException e) {
       throw new RuntimeException(e);
     }
     store.clearAll();
   }
 
   @Override
   public boolean hasMessagesForVertex(I vertexId) {
     return store.hasMessagesForVertex(vertexId);
   }
 
   @Override
   public boolean hasMessagesForPartition(int partitionId) {
     return store.hasMessagesForPartition(partitionId);
   }
 
   @Override
   public void addPartitionMessages(
       int partitionId, VertexIdMessages<I, M> messages) {
     int hash = partition2Queue.get(partitionId);
     try {
       queues[hash].put(new PartitionMessage<>(partitionId, messages));
     } catch (InterruptedException e) {
       throw new RuntimeException(e);
     }
   }
 
   @Override
   public void addMessage(I vertexId, M message) throws IOException {
     // TODO: implement if LocalBlockRunner needs async message store
     throw new UnsupportedOperationException();
   }
 
   @Override
   public void finalizeStore() {
     store.finalizeStore();
   }
 
   @Override
   public Iterable<I> getPartitionDestinationVertices(int partitionId) {
     return store.getPartitionDestinationVertices(partitionId);
   }
 
   @Override
   public void clearPartition(int partitionId) {
     store.clearPartition(partitionId);
   }
 
   @Override
   public void writePartition(DataOutput out, int partitionId)
     throws IOException {
     store.writePartition(out, partitionId);
   }
 
   @Override
   public void readFieldsForPartition(DataInput in, int partitionId)
     throws IOException {
     store.readFieldsForPartition(in, partitionId);
   }
 
   /**
    * Wait till all messages are processed and all queues are empty.
    */
   public void waitToComplete() {
     try {
       for (BlockingQueue<PartitionMessage<I, M>> queue : queues) {
         queue.put(CLEAR_QUEUE_MESSAGE);
       }
       completionSemaphore.acquire();
       completionSemaphore = new Semaphore(1 - threadsCount);
     } catch (InterruptedException e) {
       throw new RuntimeException(e);
     }
   }
 
   /**
    * This runnable has logic for background thread
    * that actually does message processing.
    */
   private class MessageStoreQueueWorker implements Runnable {
     /**
      * Queue assigned to this background thread.
      */
     private final BlockingQueue<PartitionMessage<I, M>> queue;
 
     /**
      * Constructs runnable.
      * @param queue where messages are put by client
      */
     private MessageStoreQueueWorker(
         BlockingQueue<PartitionMessage<I, M>> queue) {
       this.queue = queue;
     }
 
     @Override
     public void run() {
       PartitionMessage<I, M> message = null;
       while (true) {
         try {
           message = queue.take();
           if (message.getMessage() != null) {
             int partitionId = message.getPartitionId();
             store.addPartitionMessages(partitionId, message.getMessage());
           } else {
             completionSemaphore.release();
             if (message == SHUTDOWN_QUEUE_MESSAGE) {
               return;
             }
           }
         } catch (InterruptedException e) {
           LOG.error("MessageStoreQueueWorker.run: " + message, e);
           return;
         }
       }
     }
   }
 }