/*
    * 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.ooc.persistence;
  
  import org.apache.giraph.conf.GiraphConstants;
  import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;
  import org.apache.giraph.conf.IntConfOption;
  import org.apache.giraph.utils.ExtendedDataOutput;
  import org.apache.giraph.utils.io.BigDataInput;
  import org.apache.giraph.utils.io.BigDataOutput;
  
  import java.io.DataInput;
  import java.io.DataOutput;
  import java.io.IOException;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.LinkedBlockingDeque;
  
  /**
   * Implementation of data accessor which keeps all the data serialized but in
   * memory. Useful to keep the number of used objects under control.
   */
  public class InMemoryDataAccessor implements OutOfCoreDataAccessor {
    /** Configuration */
    private final ImmutableClassesGiraphConfiguration<?, ?, ?> conf;
    /** Factory for data outputs */
    private final PooledBigDataOutputFactory outputFactory;
    /** DataInputOutput for each DataIndex used */
    private final ConcurrentHashMap<
        DataIndex, PooledBigDataOutputFactory.Output> data;
  
    /**
     * Constructor
     *
     * @param conf Configuration
     */
    public InMemoryDataAccessor(
        ImmutableClassesGiraphConfiguration<?, ?, ?> conf) {
      this.conf = conf;
      outputFactory = new PooledBigDataOutputFactory(conf);
      data = new ConcurrentHashMap<>();
    }
  
    @Override
    public void initialize() {
      // No-op
    }
  
    @Override
    public void shutdown() {
      // No-op
    }
  
    @Override
    public int getNumAccessorThreads() {
      return GiraphConstants.NUM_OUT_OF_CORE_THREADS.get(conf);
    }
  
    @Override
    public DataInputWrapper prepareInput(int threadId,
        DataIndex index) throws IOException {
      return new InMemoryDataInputWrapper(
          new BigDataInput(data.get(index)), index);
    }
  
    @Override
    public DataOutputWrapper prepareOutput(int threadId,
        DataIndex index, boolean shouldAppend) throws IOException {
      // Don't need to worry about synchronization here since only one thread
      // can deal with one index
      PooledBigDataOutputFactory.Output output = data.get(index);
      if (output == null || !shouldAppend) {
        output = outputFactory.createOutput();
        data.put(index, output);
      }
      return new InMemoryDataOutputWrapper(output);
    }
  
    @Override
    public boolean dataExist(int threadId, DataIndex index) {
      return data.containsKey(index);
    }
  
    /**
    * {@link DataOutputWrapper} implementation for {@link InMemoryDataAccessor}
    */
   public static class InMemoryDataOutputWrapper implements DataOutputWrapper {
     /** Output to write data to */
     private final BigDataOutput output;
     /** Size of output at the moment it was created */
     private final long initialSize;
 
     /**
      * Constructor
      *
      * @param output Output to write data to
      */
     public InMemoryDataOutputWrapper(BigDataOutput output) {
       this.output = output;
       initialSize = output.getSize();
     }
 
     @Override
     public DataOutput getDataOutput() {
       return output;
     }
 
     @Override
     public long finalizeOutput() {
       return output.getSize() - initialSize;
     }
   }
 
   /**
    * {@link DataInputWrapper} implementation for {@link InMemoryDataAccessor}
    */
   public class InMemoryDataInputWrapper implements DataInputWrapper {
     /** Input to read data from */
     private final BigDataInput input;
     /** DataIndex which this wrapper belongs to */
     private final DataIndex index;
 
     /**
      * Constructor
      *
      * @param input Input to read data from
      * @param index DataIndex which this wrapper belongs to
      */
     public InMemoryDataInputWrapper(
         BigDataInput input, DataIndex index) {
       this.input = input;
       this.index = index;
     }
 
     @Override
     public DataInput getDataInput() {
       return input;
     }
 
     @Override
     public long finalizeInput(boolean deleteOnClose) {
       if (deleteOnClose) {
         data.remove(index).returnData();
       }
       return input.getPos();
     }
   }
 
   /**
    * Factory for pooled big data outputs
    */
   private static class PooledBigDataOutputFactory {
     /** How big pool of byte arrays to keep */
     public static final IntConfOption BYTE_ARRAY_POOL_SIZE =
         new IntConfOption("giraph.inMemoryDataAccessor.poolSize", 1024,
             "How big pool of byte arrays to keep");
     /** How big byte arrays to make */
     public static final IntConfOption BYTE_ARRAY_SIZE =
         new IntConfOption("giraph.inMemoryDataAccessor.byteArraySize", 1 << 21,
             "How big byte arrays to make");
 
     /** Configuration */
     private final ImmutableClassesGiraphConfiguration conf;
     /** Pool of reusable byte[] */
     private final LinkedBlockingDeque<byte[]> byteArrayPool;
     /** How big byte arrays to make */
     private final int byteArraySize;
 
     /**
      * Constructor
      *
      * @param conf Configuration
      */
     public PooledBigDataOutputFactory(
         ImmutableClassesGiraphConfiguration conf) {
       this.conf = conf;
       byteArrayPool = new LinkedBlockingDeque<>(BYTE_ARRAY_POOL_SIZE.get(conf));
       byteArraySize = BYTE_ARRAY_SIZE.get(conf);
     }
 
     /**
      * Create new output to write to
      *
      * @return Output to write to
      */
     public Output createOutput() {
       return new Output(conf);
     }
 
     /**
      * Implementation of BigDataOutput
      */
     private class Output extends BigDataOutput {
       /**
        * Constructor
        *
        * @param conf Configuration
        */
       public Output(ImmutableClassesGiraphConfiguration conf) {
         super(conf);
       }
 
       /**
        * Return all data structures related to this data output.
        * Can't use the same instance after this call anymore.
        */
       protected void returnData() {
         if (dataOutputs != null) {
           for (ExtendedDataOutput dataOutput : dataOutputs) {
             byteArrayPool.offer(dataOutput.getByteArray());
           }
         }
         byteArrayPool.offer(currentDataOutput.getByteArray());
       }
 
       @Override
       protected ExtendedDataOutput createOutput(int size) {
         byte[] data = byteArrayPool.pollLast();
         return conf.createExtendedDataOutput(
             data == null ? new byte[byteArraySize] : data, 0);
       }
 
       @Override
       protected int getMaxSize() {
         return byteArraySize;
       }
     }
   }
 }